From 98bd94aa3fc600ede05c8c92a1034cab92ee05a3 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 13 Oct 2026 17:28:23 -0700 Subject: [PATCH 01/45] Add core C++ pattern rewriter and correctness coverage Extract the engine, six example rules, API/design documentation, and correctness tests from the completed rewriter implementation without performance-evaluation artifacts. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 593 +++++ .../core/optimizer/pattern_rewriter.cc | 980 ++++++++ onnxruntime/core/optimizer/pattern_rewriter.h | 205 ++ .../core/optimizer/pattern_rewriter_rules.cc | 189 ++ .../test/optimizer/pattern_rewriter_test.cc | 1989 +++++++++++++++++ 5 files changed, 3956 insertions(+) create mode 100644 docs/Pattern_Rewriter.md create mode 100644 onnxruntime/core/optimizer/pattern_rewriter.cc create mode 100644 onnxruntime/core/optimizer/pattern_rewriter.h create mode 100644 onnxruntime/core/optimizer/pattern_rewriter_rules.cc create mode 100644 onnxruntime/test/optimizer/pattern_rewriter_test.cc diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md new file mode 100644 index 0000000000000..5fa61a9fdcc2f --- /dev/null +++ b/docs/Pattern_Rewriter.md @@ -0,0 +1,593 @@ +# Experimental indexed C++ pattern rewriter + +`onnxruntime/core/optimizer/pattern_rewriter.h` defines an opt-in `GraphTransformer`. +It is **not registered in the default optimization pipeline**, and introduces no +session configuration switch or public C ABI. C++ callers may register it using +`InferenceSession::RegisterGraphTransformer`, just like existing custom transformers. +The example predicates are provider-independent and restrict types to float32. +The engine's generic partition policy defaults to CPU or unassigned graphs; +allowing other EPs requires the caller to verify replacement kernel compatibility. + +## Source investigation and supported examples + +The ONNXScript reference revision is +[`9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36`](https://github.com/microsoft/onnxscript/tree/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36). +The following files were examined, rather than assuming that Python syntax defines +the right architecture for ORT: + +| ONNXScript source at that revision | C++ adaptation | +|---|---| +| [`rewriter/_matcher.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/_matcher.py) | Backward DAG binding, repeated-value equality, distinct matched nodes, and removal boundary checks | +| [`rewriter/_rewrite_rule.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/_rewrite_rule.py) | Separate structural patterns, semantic guards, replacement recipes, ordered priority, recursive subgraph handling | +| [`rules/common/_fuse_hardswish.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/rules/common/_fuse_hardswish.py) | `Div(Mul(Clip(Add(x,3),0,6),x),6)` and `Mul(HardSigmoid(x),x)` to HardSwish | +| [`ort_fusions/fused_matmul_rule_sets.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/ort_fusions/fused_matmul_rule_sets.py) | MatMul/scalar Div, FusedMatMul/scalar Div (preserving attributes), and last-two-dimension Transpose/MatMul | +| [`ort_fusions/rms_normalization.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/ort_fusions/rms_normalization.py) | Float32 Pow/ReduceMean/Add/Sqrt/Reciprocal/Mul/Mul to SimplifiedLayerNormalization | +| [`ort_fusions/softmax.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/ort_fusions/softmax.py) | Studied but deliberately not enabled: fp16 cast removal needs a separately verified precision/EP policy | +| [`rules/fusion/_layer_norm.py`, `LayerNormBiasFusion`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/rules/fusion/_layer_norm.py) | Test-only adaptation matches all three LayerNormalization outputs, fuses a trailing bias Add, and rejects observable statistics | + +Existing ORT `SelectorActionRegistry` already dispatches selectors by root op/domain +and version. `QuickGeluFusion` illustrates repeated inputs and provider restrictions, +while `graph_utils::FinalizeNodeFusion` illustrates edge maintenance. The new +transformer uses ORT's graph/node/edge APIs directly; it does not copy the graph to +a parallel IR. Generic backward bindings and sharing identical structural plans +are the additional capabilities over hand-written selectors. + +## Defining a rule + +The small DSL uses `Value {std::string name; int handle;}`. A name helps diagnose +patterns; only the handle determines binding identity. An input introduces a variable; +`Op` introduces one output, while `MultiOp` introduces an ordered output list. +Reusing a handle requires the identical +`NodeArg`, not merely the same shape or equivalent contents. Pattern operators +specify only name and domain. **Model-imported opset ranges belong to the rule**, +not to individual nodes or their schema since-versions. + +```cpp +using namespace onnxruntime::rewrite; +Pattern p; +auto x = p.Input("x"); +auto h = p.Op("HardSigmoid", {x}); +p.Name(h, "sigmoid"); +p.Attributes({AttributeMatcher::Equal("alpha", 1.0f / 6.0f), + AttributeMatcher::Equal("beta", 0.5f)}); +p.Op("Mul", {h, x}); +auto guard = [](const Graph&, const Match& m) { + for (const NodeArg* value : m.values) { + const auto* type = value->TypeAsProto(); + if (!type || !type->has_tensor_type() || + type->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) + return false; + } + return true; +}; +Replacement replacement("(x) => HardSwish(x)"); +Rule rule{ + "HardSigmoidMul", p, + {{kOnnxDomain, 14, 21}}, // inclusive imported opset range + guard, replacement}; +auto transformer = std::make_unique(Rules{rule}); +ORT_RETURN_IF_ERROR(session.RegisterGraphTransformer(std::move(transformer), + TransformerLevel::Level1)); +``` + +`pattern_rewriter_rules.cc` contains complete readable guard/recipe examples. +Each example defines its pattern, guard and replacement separately, assembles a +named `Rule rule`, and registers it with `rules.push_back(std::move(rule))`. +`Op(name, inputs)` defaults to the standard ONNX domain; the optional third +argument selects another domain, e.g. `p.Op("FusedMatMul", {x, y}, kMSDomain)`. +The original `Op(Operator{name, domain}, inputs)` overload remains available. +Name overloads accept literals, `std::string` and `std::string_view`; template +deduction keeps the original braced operator form unambiguous. + +`p.NamedValue("sigmoid")` retrieves a named handle. Nonempty names must be unique. +`p.Name(value, name)` updates the canonical name and existing node definitions; +subsequent operators also use that canonical name. Previously returned `Value` +copies are not renamed in place. `Value` converts to `int` for indexing +`Match::values`; an unnamed value has an empty name, and a default value has +invalid handle -1. + +In guards and replacement recipes, `m.valueOf("x")` directly returns the matched +`const NodeArg*`; there is no need to capture `p.NamedValue("x")` and index +`m.values` separately. Lookup also supports named intermediate outputs. +Unknown/empty names or unbound values throw an explicit error rather than +returning null. Names come from the current rule, even when its structural plan +is shared with a differently named pattern. Lookup scans the bounded pattern +value table and does not allocate a name map for each candidate. + +`Replacement::inputs` accepts either `Value` handles or matched `const NodeArg*` +bindings, so recipes can use `{m.valueOf("x"), m.valueOf("y")}` directly. +Pointer inputs are validated against the match's external input bindings: +intermediates, root outputs, nulls and unrelated graph values are rejected before +mutation. Matches reference immutable compiled patterns; do not retain them or +their pointers beyond a callback. + +### Attribute matchers and literal constants + +`p.Attributes(...)` attaches matchers to the most recently defined node. +Each `AttributeMatcher` contains an attribute name and a boolean predicate on +`const AttributeProto&`. Missing attributes fail the match; schema defaults are +not synthesized. Equality factories check both attribute type and value: +`Equal("axis", 1)` requires INT, `Equal("alpha", 0.5f)` requires FLOAT, +and `Equal("mode", std::string("linear"))` requires STRING. Extra unmatched +attributes are permitted. Arbitrary predicates are also supported: + +```cpp +p.Attributes({{"axis", [](const ONNX_NAMESPACE::AttributeProto& value) { + return value.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_INT && + value.i() < 0; +}}}); +``` + +Predicates are trusted, pure C++ callbacks, with the same locality and lifetime +requirements as rule guards. Duplicate/empty attribute names, empty callbacks, +or attachment before any node are errors. + +`p.Constant(3.0f)` introduces a float32 scalar initializer constraint; +`p.Constant(int64_t{3})` introduces an int64 constraint. Integer literals in C++ +select int64; `double` arguments are converted to float32. Float constants must +be finite. Matching requires an immutable rank-zero initializer of the exact +type and value; a one-element vector, an overridable initializer, or external-data +initializer does not match. These leaf handles can be used as replacement inputs, +but matching never creates tensors. + +### Textual patterns + +The explicit `Pattern(std::string_view text)` constructor parses a textual pattern. +The default constructor remains available for the programmatic builder. +Malformed text throws `OnnxRuntimeException` with a byte-offset diagnostic +(or aborts in exception-disabled builds, following ORT's throwing macros). +The requested branched pattern can be written directly: + +```cpp +Pattern p(R"( + (x, low, high, divisor) => + add = Add(x, 3.0) + clip = Clip(add, low, high) + mul = Mul(clip, x) + return Div(mul, divisor) +)"); +``` + +The header declares named inputs. Assignment names label operator outputs. +Reusing `x` enforces exact binding identity. Decimal/scientific literals become +float32 scalar constants, while integer spellings become int64 constants. +This text describes structure, **not a semantically complete HardSwish rule**: +the rule guard must additionally establish float32, low=0, high=6 and divisor=6. +The numerical test retains those checks before executing the replacement. + +The actual example rule eliminates those fixed bindings entirely: + +```cpp +Pattern p(R"( + (x) => + add = Add(x, 3.0) + clip = Clip(add, 0.0, 6.0) + mul = Mul(clip, x) + return Div(mul, 6.0) +)"); +``` + +Its guard only checks float32, since the pattern checks every fixed scalar. +RMSNorm similarly uses `Pow(x, 2.0)` directly; epsilon remains a named binding +because its value varies and is copied to the fused operator's attributes. +MatMul division factors also remain variable for computed scaling attributes. + +Attributes use Python keyword syntax, after positional inputs: +`h = HardSigmoid(x, alpha = 0.16666667, beta = 0.5)`. +Signed/scientific numbers and single/double-quoted string attributes are supported. +Only escaped quotes and backslashes are supported in strings. +The parsed spelling determines the attribute type: `axis=1` differs from `axis=1.0`. +An empty quoted string is valid. + +Multiple-output assignments and qualified domains are supported: + +```text +(x) => a, b = Split(x, axis=0); return Concat(a, b, axis=0) +(x, y) => mm = com.microsoft::FusedMatMul(x, y, alpha=0.5); return Identity(mm) +``` + +Unqualified operators use the standard ONNX domain. Whitespace/newlines separate +tokens; semicolons may separate assignments. Values are ASCII identifiers. +Definitions must precede use; every input and node must reach the final operator. +There must be a final `return Op(...)`, with no subsequent statements. +All six example rules use the string constructor and retrieve bindings with +`Match::valueOf` for their guards and dynamic replacement recipes; static +replacement plans bind inputs once during compilation. Tests use the constructor for parsed +patterns and test-only Identity, Split/Concat and LayerNorm-bias rules; builder +tests and multi-output-root definitions retain the programmatic API. + +The parser checks duplicate names/attributes, undefined references, numeric range, +disconnected definitions and the same 64-node/256-value/64-input-output limits as +the C++ builder. Errors include a byte offset; `p = Pattern(text)` leaves `p` +unchanged if construction fails. +A one-node parsed pattern is valid, but cannot form a node-count-decreasing rule. + +The language intentionally excludes nested calls, lists/tensor literals, optional +placeholders, booleans, comments, `return` of an existing value and multi-output +`return` calls (use the C++ API for multi-output roots). It describes only matching: +opset requirements and semantic guards remain explicit C++. + +### Reusable replacement plans + +`Replacement(std::string_view)` parses a reusable replacement definition, once: + +```cpp +Replacement replacement("(x) => HardSwish(x)"); +Rule rule{"HardSwish", pattern, {{kOnnxDomain, 14, 21}}, guard, replacement}; +``` + +The final operator call may omit `return`. Multi-node definitions use the same +assignment, domain and keyword-attribute syntax as patterns: + +```cpp +Replacement replacement(R"( + (x) => + h = HardSigmoid(x, alpha=0.16666667, beta=0.5) + return Mul(h, x) +)"); +``` + +This defines a two-node replacement, not a callback that parses a string on every +match. The constructor builds an immutable plan containing operator/domain names, +literal attributes, external input names and indexed earlier-output references. +Copies share that plan. When the plan is supplied directly as `Rule::replacement`, +transformer construction binds its external names to the rule's input handles, +and checks declared domains and node-count reduction. Binding is per rule, so +reusing a plan across differently ordered input headers is safe. + +At runtime the callback returns a shared plan reference. There is no parsing, +attribute-literal construction, or name lookup for these prebound inputs. +The engine binds the external `NodeArg`s, allocates fresh intermediate values, +creates nodes in definition order, and applies the stored attributes. +Callbacks that need computed attributes can use extra bindings with a captured +parsed plan (below), or return programmatic recipes. Do not construct +`Replacement(text)` inside a per-match callback. + +Replacement headers may select a subset of the matched pattern's inputs. +Tensor parameters cannot refer to matched intermediates. All declared replacement +parameters must be used as tensor inputs or attribute bindings, and all nodes +must reach the final operator. A multi-output assignment such as +`a, b = Split(x, axis=0)` declares the exact internal arity. +The final operator inherits **all matched root outputs in order**, including a +multi-output root; its schema must accept that count. Intermediate outputs cannot +be mapped independently to observable root slots yet. + +Numeric positional literals are synthesized as immutable rank-zero float32 or +int64 initializers, using the same literal typing as patterns: +`Replacement("(x) => Add(x, 0.0)")`. No Constant kernel or subsequent +constant-folding pass is required. There is no general tensor/list synthesis. + +Programmatic sequences use `NodeRecipe` and `Output{node_index, output_slot}`: + +```cpp +Replacement replacement(InlinedVector{ + {"HardSigmoid", kOnnxDomain, {std::string("x")}, sigmoid_attributes, 1}, + {"Mul", kOnnxDomain, {Replacement::Output{0, 0}, std::string("x")}, {}, 1}}); +``` + +`NodeRecipe` contains op, domain, inputs, attributes and internal output count. +It also has `attribute_bindings`, a map from target attribute names to extra +binding names. +Its inputs may be names, existing `Value` handles, matched `NodeArg*` bindings, +earlier-output references or `Replacement::Scalar{...}` literals. +The final node's output count is inherited from the root rather than this field. +The original four-argument single-node constructor remains supported, including +computed-attribute callbacks. Its public single-node fields cannot be combined +with an immutable parsed/programmatic plan; mixed representations are rejected. + +#### Computed attributes through extra bindings + +Replacement keyword values may refer to names declared in its header. +Tensor arguments come from the pattern match; extra attribute arguments come +from `Match::update`, which accepts a concrete `AttributeProto`: + +```cpp +Replacement plan("(x, y, alpha) => com.microsoft.FusedMatMul(x, y, alpha=alpha)"); +auto replacement = [plan](const Graph& graph, const Match& m) { + m.update("alpha", FloatAttribute("alpha", + 1.0f / *Scalar(graph, *m.valueOf("divisor")))); + return plan; +}; +Rule rule{"MatMulDiv", pattern, opsets, guard, replacement}; +``` + +Here `FloatAttribute` and `Scalar` are the helpers used by the example rules; +the guard establishes float32 and a finite, nonzero scalar divisor. +The plan is parsed **once when defining the rule**, not once per successful +match. Capturing/returning it shares the immutable plan. Computed attributes +are looked up and copied into the new nodes at rewrite time. Captured callback +plans still resolve tensor names at runtime; only directly supplied plans have +their tensor names prebound by transformer construction. + +`com.microsoft.FusedMatMul` is accepted in replacement text, as is the existing +`com.microsoft::FusedMatMul` spelling. Literal attributes and bound attributes +can be mixed, e.g. `axis=-1, epsilon=fused_epsilon, stash_type=1`. +Multiple nodes can share a binding, and target attribute names can differ from +binding names. Attribute proto names are normalized to the target names. +The proto carries the type, including list/tensor/string types not representable +as text literals; imported operator schemas validate the resolved attributes. + +Extra bindings are invocation-local state, separate from immutable matched +graph values. `update` is callable on the callback's `const Match&`, because it +changes only this extra state, not the graph or matched values. Repeated updates +to an extra name replace its value. Guards must remain pure; the engine clears +extra state between candidates and before invoking the winning recipe. No binding +is retained between rewrites, graph runs or rules; match-only mode never invokes +recipes. `attributeOf(name)` retrieves a previously set extra attribute. + +Empty names, undefined attribute types, reference attributes, missing bindings +and attempts to shadow any named pattern value are explicit errors. +A header name cannot be used both as a tensor argument and an attribute binding; +attribute references must name header parameters, not intermediate outputs. +All required extras are checked before graph mutation. Schema validation failures +use the same staged-node rollback as literal attributes. As with other recipes, +callbacks must establish semantic/type compatibility; extra bindings are not +an automatic semantic guard or a general graph transaction. + +The MatMulDiv and RMSNorm examples use captured plans with computed bindings. +TransposeMatMul uses a literal parsed replacement. FusedMatMulDiv also captures +a parsed plan, computing `alpha` and binding `transA`, `transB`, `transBatchA` +and `transBatchB` from the original node. Missing transpose attributes use the +imported schema's defaults; the replacement emits all five version-1 attributes +explicitly, preserving their effective values. All six examples therefore use +parsed replacements. Extra bindings describe individual attributes, not an +attribute-map spread; a rule for a schema with other attributes must explicitly +preserve those too. + +Every replacement must contain at least one node and **strictly fewer nodes +than the matched pattern**. Equal-size rewrites and expansions are rejected, +including dynamic recipes, retaining the node-count termination policy. +All internal references must point backward to valid output slots; disconnected +nodes and undeclared domains are errors. Limits remain 64 nodes and 64 inputs/ +outputs per node. Pattern parsing additionally bounds named values to 256. + +All replacement nodes are staged and their imported schemas/attributes checked +before deleting matched nodes or rewiring edges. A schema failure removes every +staged node and synthesized initializer, leaving the original live graph intact; +unused staging `NodeArg`s and consumed unique-name IDs may remain. Final Resolve +is still responsible for type/shape inference and validation; arbitrary computed +recipes must establish semantic/type compatibility, and post-commit Resolve +failure is not a general graph transaction rollback. +All added nodes inherit root provider/layer metadata. Producer/consumer maps, +internal edges and external root edges are updated; every added node seeds the +invalidation worklist. Replacement notification attributes the removed pattern +to the final node. + +### Multiple outputs + +`MultiOp(name, inputs, output_count, domain)` returns an ordered vector of distinct +value handles. `domain` is optional, and an `Operator` overload is also available. +For example: + +```cpp +Pattern p; +auto x = p.Input(); +auto pieces = p.MultiOp("Split", {x}, 2); +p.Op("Concat", {pieces[0], pieces[1]}); +``` + +The output count must match the node exactly. Matching can find a producer +through any output, including its second or later output, and checks every +bound output against its exact output slot. Unused outputs of a connected node +are allowed in the definition, but cannot be graph outputs or have external +consumers if the node is deleted. The tests include Split/Concat and a genuine +ONNXScript LayerNorm-bias adaptation with unused mean/inverse-standard-deviation +outputs; exposing either statistic prevents its deletion. + +A multi-output root is also supported. The final replacement node inherits +**all root outputs in the original order**, including their names, types, shapes, +graph-output status, and consumers. The recipe does not need an output list. +Its schema must permit that output count. Identity/Split to Split tests execute +both root outputs and downstream fanout, including repeated uses of output 1. +This does not yet support mapping observable outputs from multiple matched nodes +to arbitrary replacement slots, or matching omitted optional output placeholders. + +### Guards, applicability and integration + +Guards and recipes are trusted C++ code: they must be pure, guards must depend on +local bindings/attributes/constants rather than mutable global graph conditions, +they must not retain match +pointers, and must establish semantic equivalence and replacement type +compatibility. The caller must establish kernel compatibility for its EP policy. +The engine provides structural and deletion safety; it cannot prove +an arbitrary callback's mathematical correctness. Guard false means an ordinary +nonmatch; invalid patterns or recipes report explicit errors. + +Each `OpsetRequirement` is `{domain, min_version, max_version}`, with inclusive, +positive bounds. A rule must declare every domain used by its pattern or +replacement; additional domain requirements are allowed. All requirements must +be satisfied by the graph's import map. Missing imports or out-of-range versions +make the whole rule inapplicable, before matching or guards. Duplicate domains +(including `""`/`"ai.onnx"` aliases), reversed ranges and undeclared pattern +domains are rejected when compiling rules. + +For example, MatMul/Div to FusedMatMul declares +`{{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}`. The standard-domain nodes use the +model's ONNX import; the replacement uses the existing Microsoft-domain import. +Rule authors are responsible for equivalence throughout the declared ranges. +The engine neither derives nor compares node schema since-versions. Replacement +schemas are resolved normally from imports and validated before deletion. + +`Apply` recurses into subgraphs and resolves the graph through the normal ORT +transformer wrapper. For per-call statistics and explicit graph resolution: + +```cpp +PatternRewriter rewriter(MakeExampleRules()); +Statistics stats; +ORT_RETURN_IF_ERROR(rewriter.Run(graph, stats)); // already-resolved graph, one graph only +ORT_RETURN_IF_ERROR(graph.Resolve()); // required after rewriting +``` + +Callers must not concurrently mutate a graph. Compiled rules and options are +immutable and per-call state/statistics are local. + +### Generic execution-provider policy + +Rules do not inspect EP assignments. After structural matching and before rule +guards, the engine applies `CanRewritePartition`: all matched nodes must have +the same assignment, and an assigned EP must belong to +`Options::compatible_execution_providers`. Unassigned nodes are always eligible. +The default allowed set contains only CPU, preserving the verified execution scope. +Callers can opt into an EP (or clear the set for unassigned-only rewrites): + +```cpp +Options options; +options.compatible_execution_providers.insert(kCudaExecutionProvider); +PatternRewriter rewriter(MakeExampleRules(), options); +``` + +The replacement inherits the root's assignment. Allowing an EP is a caller assertion +of kernel/precision compatibility, not automatic kernel discovery or evidence of +GPU execution. Mixed assignments are rejected even if both EPs are allowed. This +generic policy also applies in match-only mode, which measures safe rewrite +eligibility rather than partition-independent structural matches. + +## Compilation, matching, mutation, and termination + +Construction validates connected acyclic definitions, used inputs, rule-level +domain/opset requirements, and limits of 64 nodes, 256 values and 64 inputs/outputs per operator. +It buckets groups by root op/domain and interns identical structural signatures, +including ordered output handles, binding identities and scalar literal type/value +constraints, but not diagnostic names or rule-level import ranges. +Rules with identical topology and literal constraints can share one structural +match even when their guards or applicability ranges differ. Patterns containing +attribute predicates currently receive separate plans, including equality predicates: +opaque `std::function` callbacks cannot safely be compared. This conservative +choice preserves correctness at the cost of repeated matching for attribute-heavy +rule sets; arbitrary callbacks also have no engine-imposed time bound. +There is **no arbitrary-prefix trie**: different structural +plans still match independently. This is intentionally a simpler, measurable +starting point than a decision DAG compiler. + +Each call filters rules once against `Graph::DomainToVersionMap()`, retaining +registration order within each structural group. No applicability state is cached +across graphs; recursive subgraph calls use their own graph import context. +Filtering costs the sum of rule requirement counts, outside the node traversal. +If no rules apply, the call returns without traversing nodes. Model-local function +rewriting remains unsupported; future function rewriting must use that function's +imports, not the enclosing model's. + +The call then builds a producer index and topological worklist once. Matching walks +the compiled node definitions backward without recursive search, alternation, +commutative backtracking, or unbounded neighborhood exploration. Each node checks +op/domain, exact input/output counts, repeated bindings and layer annotation +agreement, plus any attribute predicates and scalar literal constraints. +Provider eligibility is a separate generic check. All non-root outputs +must be internal: graph outputs, +external consumers (including implicit subgraph captures), and pattern inputs +produced inside the deletion set prevent replacement. + +Registration order defines priority even across different structural groups. +Match-only mode reports all passing rule guards; rewrite mode selects only the +first passing rule. Replacements may consume **only matched input variables, +new scalar literals or earlier replacement outputs**, retain the +root's output `NodeArg`s (name/type/shape and external users), and inherit its +execution provider and layer annotation. The replacement domain must be declared +in the rule requirements and already imported for every added node. Replacement schema/attribute checks +precede deletion. ORT edges and producer/consumer indexes are updated; removed +nodes are reported via `NotifyNodeReplacement`. Final Resolve validates types and +inferred shapes, rather than silently ignoring graph-validation errors. + +After mutation, all replacement nodes and surviving input producers seed a bounded +forward invalidation walk. Descendant roots up to the maximum pattern size are +requeued, since producer changes or reduced use counts can unlock deeper patterns. +Queue deduplication avoids redundant pending entries; removed IDs are skipped. +Traversal and rule priority are deterministic. + +Every rule replaces at least two nodes with a smaller nonempty node sequence. Thus every +successful rewrite strictly decreases live node count, and a graph with N nodes +admits at most N-1 rewrites. This is a termination policy, **not a runtime bound**: +guards, fanout, structural checks, and worklist invalidation still cost time. +Expansion rules and one-node rewrites are deliberately rejected. + +For match-only workloads, let U be the number of distinct structural plans in a +root bucket and P the maximum pattern size. Structure work is bounded by the +sum over visited roots of U times pattern/boundary checking cost, plus passing +plan guard costs. Boundary checks include matched-node fanout; they are not +constant-time. The reference mode performs the same checks per rule instead, +with cheap root op/domain rejection before allocating binding buffers. +Rewrite workloads additionally pay edge updates and bounded descendant walks; +large fanout can make these expensive. Arbitrary trusted guards have no imposed +time bound. These structural bounds do not provide mathematical worst-case runtime +guarantees for arbitrary rules or models. + +`visited` counts live worklist roots, including revisits. `candidates` counts +enumerated rule IDs (all rules in the reference mode, root-indexed rules in the +indexed mode); `structural_checks` counts plan attempts, including cheap root +rejections in the reference. `guards` counts semantic callback evaluations after +structural and deletion-boundary checks. `matches` counts all passing guards in +match-only mode, or selected winners in rewrite mode; `rewrites` counts completed +mutations. Match-only therefore measures safe rewrite eligibility, not a search +that ignores external consumers. + +## Intentional limitations and precision + +Matched and replacement nodes may have multiple outputs, but the final replacement +node must preserve the root's complete ordered output list. Optional/missing +inputs or outputs, variadic patterns, +model-local function rewriting, general tensor synthesis, domain-import edits, alternative +patterns and implicit commutation are unsupported. Control-flow nodes are not +matched, but `Apply` rewrites inside their subgraphs. External tensor data and +overridable initializers do not satisfy example constant guards; scalar constants +must be finite float32 rank-zero initializers. Singleton vectors are deliberately +not treated as scalars, since broadcasting can otherwise change the output rank. +Constant nodes require an earlier constant-folding pass. + +The examples support a conservative subset, not every ONNXScript variant: +RMSNorm requires the last axis, keepdims, positive finite epsilon, float32, +known positive-width matching one-dimensional scale, and pre-opset-18 attribute +axes. Transpose requires known rank >=2 and exactly the last-two-axis swap. +RMSNorm declares ONNX opsets 7 through 17. MatMul Div requires a finite nonzero +scalar and finite reciprocal. The FusedMatMul examples declare ONNX opsets 9 +through 21 and require an existing `com.microsoft` opset-1 import; they do not add +one. The upper bounds are conservative supported ranges, not inferred schema +version boundaries; extending them requires reviewing and testing semantics. +HardSwish requires exact float32 constants/attributes and model opset 14 through +21. The tested CPU build executes HardSwish by expanding its ONNX function body +to HardSigmoid/Mul; RMSNorm and FusedMatMul execute native CPU kernels. +Floating-point fusions change evaluation order; equivalence means tested +tolerance, not bitwise identity. NaNs/infinities and every extreme finite input +are not claimed to have bitwise-equivalent behavior. + +## Correctness tests + +The normal tests live in `onnxruntime/test/optimizer/pattern_rewriter_test.cc`. +They include executed CPU equivalence and structural assertions, not only match +counts. Use an isolated build directory and virtual environment: + +```powershell +python -m venv .venv +.\.venv\Scripts\python.exe tools\ci_build\build.py --build_dir build\rewriter --config Release --update --parallel 4 --skip_tests --cmake_generator "Visual Studio 17 2022" --cmake_extra_defines FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER +.\.venv\Scripts\python.exe tools\ci_build\build.py --build_dir build\rewriter --config Release --build --parallel 4 --target onnxruntime_test_all --skip_tests +Set-Location build\rewriter\Release\Release +.\onnxruntime_test_all.exe --gtest_filter="PatternRewriterTests.*" +``` + +Ensure CMake is on PATH (for example, Visual Studio's bundled CMake). +`FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER` avoids accidentally importing an +incompatible installed dependency, such as a Debug-only protobuf installation +while building Release. Dependencies are built inside this isolated +build tree instead of modifying the shared installation. + +## Decoder correctness topology + +`BuildDecoder` is a **hand-constructed, model-derived eager Llama-style topology**, +not an exported pretrained model or a claim of identical exporter node order. +Sources at ORT base revision +[`42518d89d7c5393baf97d69f8c2670668c40decf`](https://github.com/microsoft/onnxruntime/tree/42518d89d7c5393baf97d69f8c2670668c40decf) +are `onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py` +(eager decoder, dynamic past/present KV exports), +`models/llama/llama_inputs.py` (cache layout), +`onnxruntime/test/optimizer/group_query_attention_pre_norm_fusion_test.cc` +(small Q/K/V/head/cache geometry), and ONNXScript's RMSNorm source above. + +Each block has two decomposed RMSNorms, Q/K/V projections, reshape/transposes, +split-half RoPE via Slice/Neg/Concat and gathered sine/cosine position caches, +past/current KV concatenation with per-layer graph outputs, grouped-head expansion, +QK transpose/MatMul/scaling/causal-mask/Softmax/AV attention, output projection, +two residual adds, and a SiLU-gated up/gate/down MLP. B=1, query heads=4, +KV heads=2, head size=4, hidden width=16, MLP width=64, past length=3. +Sine/cosine caches use standard paired frequencies; weights are tiny random +float32 tensors. No tokenizer, embedding lookup, LM head, dynamic position/cache +control flow, mixed precision, quantization or production tensor sizes are modeled. +Both sequence=1 decode and sequence=4 causal prefill have executed equivalence tests. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc new file mode 100644 index 0000000000000..30ffc78bfc0c8 --- /dev/null +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -0,0 +1,980 @@ +// 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 "core/graph/graph_utils.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 << pattern.ValueCount() << ':'; + 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([&](auto value) { stream << value; }, 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 MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, const Producers& producers, + Match& match) { + 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 false; + } + match.values.assign(pattern.ValueCount(), nullptr); + match.nodes.assign(pattern.Nodes().size(), nullptr); + match.nodes.back() = &root; + 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 || !distinct.insert(node->Index()).second || node->OpType() != expected.op.name || + node->Domain() != expected.op.domain || node->InputDefs().size() != expected.inputs.size() || + node->OutputDefs().size() != expected.outputs.size() || !node->GetAttributeNameToSubgraphMap().empty() || + !node->ImplicitInputDefs().empty() || node->GetLayeringAnnotation() != root.GetLayeringAnnotation()) { + return false; + } + match.nodes[i] = node; + for (const auto& attribute : expected.attributes) { + auto found = node->GetAttributes().find(attribute.name); + if (found == node->GetAttributes().end() || !attribute.Matches(found->second)) return false; + } + 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 false; + binding = output; + } + for (size_t j = 0; j < expected.inputs.size(); ++j) { + const NodeArg* actual = node->InputDefs()[j]; + const NodeArg*& binding = match.values[expected.inputs[j]]; + if (!actual->Exists() || (binding && binding != actual)) return false; + binding = actual; + } + } + // Input variables must be outside the deletion set, even when aliasing another binding. + 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 false; + } + for (const auto& [handle, expected] : pattern.Constants()) { + const auto* tensor = graph_utils::GetConstantInitializer(graph, match.values[handle]->Name()); + if (!tensor || tensor->dims_size() != 0 || + tensor->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL) + return false; + const bool floating = std::holds_alternative(expected); + if (tensor->data_type() != + (floating ? ONNX_NAMESPACE::TensorProto_DataType_FLOAT : ONNX_NAMESPACE::TensorProto_DataType_INT64)) + return false; + Initializer initializer(graph, *tensor, graph.ModelPath()); + if (initializer.size() != 1) return false; + if (floating ? *initializer.data() != std::get(expected) + : *initializer.data() != std::get(expected)) + return false; + } + for (size_t i = 0; i + 1 < match.nodes.size(); ++i) { + const Node& node = *match.nodes[i]; + if (graph.NodeProducesGraphOutput(node)) return false; + for (auto edge = node.OutputEdgesBegin(); edge != node.OutputEdgesEnd(); ++edge) { + if (!distinct.count(edge->GetNode().Index())) return false; + } + } + 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; }); +} + +Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const Replacement& recipe, + Producers& producers, InlinedVector& added) { + ORT_RETURN_IF(recipe.Plan() && (!recipe.op.empty() || !recipe.domain.empty() || !recipe.inputs.empty() || + !recipe.attributes.empty()), + "A replacement plan cannot be combined with single-node recipe fields."); + InlinedVector single; + if (!recipe.Plan()) single.push_back({recipe.op, recipe.domain, recipe.inputs, recipe.attributes, 1}); + const auto& nodes = recipe.Plan() ? *recipe.Plan() : single; + ORT_RETURN_IF(nodes.empty() || nodes.size() >= match.nodes.size(), + "Replacements must contain fewer nodes than the matched pattern."); + const Node& root = *match.nodes.back(); + InlinedVector> external_inputs; + for (size_t i = 0; i < nodes.size(); ++i) { + const auto& node = nodes[i]; + for (const auto& [name, binding] : node.attribute_bindings) { + ORT_RETURN_IF(name.empty() || node.attributes.count(name), "Invalid or duplicate replacement attribute: ", name); + match.attributeOf(binding); + } + 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 (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; + } + Value value; + if (const auto* handle = std::get_if(&input)) { + value = *handle; + } else { + const auto* argument = std::holds_alternative(input) ? match.valueOf(std::get(input)) + : std::get(input); + for (int id = 0; id < static_cast(match.values.size()); ++id) { + if (argument && match.values[id] == argument && pattern.IsInput(id)) { + value = id; + break; + } + } + } + ORT_RETURN_IF( + value.handle < 0 || static_cast(value.handle) >= match.values.size() || !pattern.IsInput(value), + "Replacement inputs must be pattern input variables."); + inputs.push_back(graph.GetNodeArg(match.values[value]->Name())); + } + external_inputs.push_back(std::move(inputs)); + } + // Require a connected sequence ending at the node that inherits the root outputs. + 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); + } + InlinedVector initializers; + auto rollback = [&]() { + for (auto it = added.rbegin(); it != added.rend(); ++it) graph.RemoveNode(*it); + for (const auto& name : initializers) graph.RemoveInitializedTensor(name); + added.clear(); + }; + // 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")); + if (const auto* floating = std::get_if(&scalar->value)) { + tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor.add_float_data(*floating); + } else { + tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + tensor.add_int64_data(std::get(scalar->value)); + } + ONNX_NAMESPACE::TypeProto type; + type.mutable_tensor_type()->set_elem_type(tensor.data_type()); + type.mutable_tensor_type()->mutable_shape(); + inputs[j] = &graph.GetOrCreateNodeArg(tensor.name(), &type); + initializers.push_back(tensor.name()); + graph.AddInitializedTensor(tensor); + } + } + InlinedVector outputs; + if (i + 1 == nodes.size()) { + for (const NodeArg* output : root.OutputDefs()) outputs.push_back(graph.GetNodeArg(output->Name())); + } else { + for (size_t j = 0; j < nodes[i].output_count; ++j) + outputs.push_back(&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("rewrite_value"), nullptr)); + } + const auto& node = nodes[i]; + const auto& domain = node.domain == kOnnxDomainAlias ? kOnnxDomain : node.domain; + NodeAttributes resolved_attributes; + const NodeAttributes* attributes = &node.attributes; + if (!node.attribute_bindings.empty()) { + resolved_attributes = node.attributes; + for (const auto& [name, binding] : node.attribute_bindings) { + auto value = match.attributeOf(binding); + value.set_name(name); + resolved_attributes.emplace(name, std::move(value)); + } + attributes = &resolved_attributes; + } + Node& replacement = graph.AddNode(graph.GenerateNodeName("rewrite"), node.op, "PatternRewriter", inputs, outputs, + root, attributes, domain); + added.push_back(replacement.Index()); + if (!graph.SetOpSchemaFromRegistryForNode(replacement)) { + rollback(); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Replacement schema is unavailable: ", node.op); + } + ORT_TRY { + ONNX_NAMESPACE::NodeProto proto; + replacement.ToProto(proto); + replacement.Op()->Verify(proto); + } + ORT_CATCH(const ONNX_NAMESPACE::checker::ValidationError& error) { + rollback(); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid replacement: ", error.what()); + } + replacement.SetExecutionProviderType(root.GetExecutionProviderType()); + } + struct Edge { + NodeIndex destination; + int source_arg; + int destination_arg; + }; + InlinedVector outgoing; + for (auto it = root.OutputEdgesBegin(); it != root.OutputEdgesEnd(); ++it) { + outgoing.push_back({it->GetNode().Index(), it->GetSrcArgIndex(), 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()) { + 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]; + 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) + graph.AddEdge(added.back(), edge.destination, edge.source_arg, 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::NamedValue(std::string_view name) const { + 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(std::initializer_list matchers) { + Attributes(InlinedVector(matchers)); +} + +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())); +} + +Value Pattern::Op(Operator op, std::initializer_list inputs) { return MultiOp(std::move(op), inputs, 1)[0]; } + +InlinedVector Pattern::MultiOp(Operator op, std::initializer_list inputs, size_t output_count) { + return MultiOp(std::move(op), gsl::span(inputs.begin(), inputs.size()), output_count); +} + +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 >= 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(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 { +class PatternParser { + public: + explicit PatternParser(std::string_view text, InlinedVector* attributes = nullptr, + InlinedVector, 1>* bindings = nullptr) + : text_(text), replacement_attributes_(attributes), replacement_bindings_(bindings) {} + + 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)); + ORT_RETURN_IF(symbols_.count(name), Error("duplicate input name")); + ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("too many values")); + symbols_.emplace(name, pattern_.Input(name)); + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + } + ORT_RETURN_IF_ERROR(Expect(")")); + ORT_RETURN_IF_ERROR(Expect("=>")); + bool returned = false; + while (kind_ != Kind::End) { + if (kind_ == Kind::Symbol && token_ == ";") { + ORT_RETURN_IF_ERROR(Next()); + continue; + } + InlinedVector names; + if (kind_ == Kind::Identifier && token_ == "return") { + returned = true; + ORT_RETURN_IF_ERROR(Next()); + } else { + const size_t start = token_start_; + do { + std::string name; + ORT_RETURN_IF_ERROR(Identifier(name, replacement_attributes_ && names.empty())); + ORT_RETURN_IF(std::find(names.begin(), names.end(), name) != names.end(), Error("duplicate assignment name")); + names.push_back(std::move(name)); + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + if (replacement_attributes_ && names.size() == 1 && (token_ == "(" || token_ == "::")) { + position_ = start; + ORT_RETURN_IF_ERROR(Next()); + names.clear(); + returned = true; + } else { + for (const auto& name : names) ORT_RETURN_IF(symbols_.count(name), Error("duplicate assignment name")); + ORT_RETURN_IF(std::any_of(names.begin(), names.end(), + [](const auto& name) { return name.find('.') != std::string::npos; }), + Error("expected value identifier")); + ORT_RETURN_IF_ERROR(Expect("=")); + } + } + ORT_RETURN_IF_ERROR(Call(names)); + if (returned) { + if (kind_ == Kind::Symbol && token_ == ";") ORT_RETURN_IF_ERROR(Next()); + ORT_RETURN_IF(kind_ != Kind::End, Error("unexpected text after return")); + break; + } + } + ORT_RETURN_IF(!returned, Error("missing return")); + InlinedHashSet reachable; + for (Value output : pattern_.Nodes().back().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) reachable.insert(input.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() && std::isspace(static_cast(text_[position_]))) ++position_; + token_start_ = position_; + token_.clear(); + kind_ = Kind::End; + if (position_ == text_.size()) return Status::OK(); + kind_ = Kind::Symbol; + char c = text_[position_++]; + if (Alpha(c)) { + kind_ = Kind::Identifier; + while (position_ < text_.size() && + (Alpha(text_[position_]) || Digit(text_[position_]) || 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" || (!allow_domain && token_.find('.') != std::string::npos), + Error("expected value identifier")); + name = token_; + return Next(); + } + + 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; + 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() || !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 Call(const InlinedVector& names) { + 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 if (replacement_attributes_) { + 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")); + } + } + 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(!attribute_names.insert(name).second, Error("duplicate attribute")); + ORT_RETURN_IF_ERROR(Next()); + if (replacement_bindings_ && kind_ == Kind::Identifier) { + auto found = symbols_.find(token_); + ORT_RETURN_IF(found == symbols_.end() || !pattern_.IsInput(found->second), + Error("attribute binding must reference a declared replacement input")); + bindings.emplace(name, token_); + attribute_inputs_.insert(found->second.handle); + ORT_RETURN_IF_ERROR(Next()); + } 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 { + 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")); + 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")); + std::variant number; + ORT_RETURN_IF_ERROR(Number(number)); + ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("too many values")); + 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); + 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) { + 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_; +}; +} // namespace + +Pattern::Pattern(std::string_view text) { ORT_THROW_IF_ERROR(PatternParser(text).Parse(*this)); } + +Replacement::Replacement(InlinedVector nodes) { + ORT_ENFORCE(!nodes.empty() && nodes.size() <= 64, "Replacements need between 1 and 64 nodes."); + plan_ = std::make_shared>(std::move(nodes)); +} + +Replacement::Replacement(std::string_view text) { + Pattern parsed; + InlinedVector parsed_attributes; + InlinedVector, 1> parsed_bindings; + ORT_THROW_IF_ERROR(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) { + 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)); + } + 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."); + auto value = std::find_if(pattern->values_.begin(), pattern->values_.end(), + [&](const Value& candidate) { return candidate.name == name; }); + ORT_ENFORCE(!name.empty() && value != pattern->values_.end(), "Unknown pattern value: ", name); + ORT_ENFORCE(value->handle >= 0 && static_cast(value->handle) < values.size() && values[value->handle], + "Pattern value is not bound: ", name); + return values[value->handle]; +} + +void Match::update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) const { + 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."); + if (pattern) { + 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); + } + 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 { + auto found = extra_attributes_.find(std::string(name)); + ORT_ENFORCE(!name.empty() && found != extra_attributes_.end(), "Missing replacement attribute binding: ", name); + return found->second; +} + +PatternRewriter::PatternRewriter(Rules rules, Options options) + : GraphTransformer("PatternRewriter"), rules_(std::move(rules)), options_(options) { + InlinedHashMap signatures; + for (size_t i = 0; i < rules_.size(); ++i) { + Rule& rule = rules_[i]; + ORT_ENFORCE(!rule.name.empty() && rule.pattern.Nodes().size() >= 2 && rule.guard && rule.replacement, + "Rules need a name, at least two nodes, a guard 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_; + ORT_ENFORCE(nodes.size() < rule.pattern.Nodes().size(), + "Replacements must contain fewer nodes than the matched pattern."); + 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) { + ORT_ENFORCE(std::none_of(rule.pattern.values_.begin(), rule.pattern.values_.end(), + [&](const Value& value) { return value.name == binding; }), + "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), "Replacement inputs must be pattern 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) reachable.insert(input.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(); })) + 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."); + // 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); + } + 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; + } + } + // Nodes are revisited only after a strictly node-count-decreasing rewrite. + for (size_t cursor = 0; cursor < worklist.size(); ++cursor) { + queued.erase(worklist[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; + auto consider = [&](const Pattern& pattern, gsl::span rule_ids) { + statistics.candidates += rule_ids.size(); + ++statistics.structural_checks; + Match match; + if (!MatchPattern(pattern, graph, *root, producers, match) || !CanRewritePartition(match, options_)) return; + for (size_t id : rule_ids) { + if (!options_.match_only && id >= winner) break; + match.pattern = &rules_[id].pattern; + match.ClearUpdates(); + ++statistics.guards; + if (!rules_[id].guard(graph, match)) continue; + if (options_.match_only) ++statistics.matches; + if (id < winner) { + winner = id; + winning_match = match; + winning_pattern = &pattern; + } + if (!options_.match_only) break; + } + }; + if (options_.indexed) { + auto it = roots_.find(Key(root->OpType(), root->Domain())); + if (it != roots_.end()) { + for (size_t id : it->second) { + const auto& group = groups_[id]; + const auto& ids = applicable_groups[id]; + if (ids.empty() || (!options_.match_only && ids.front() >= winner)) continue; + consider(group.pattern, ids); + } + } + } else { + for (size_t id : applicable_rules) { + consider(rules_[id].pattern, gsl::span(&id, 1)); + if (!options_.match_only && winner != std::numeric_limits::max()) 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; + winning_match.ClearUpdates(); + 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; }); + }; + if (recipe.Plan()) { + for (const auto& node : *recipe.Plan()) + ORT_RETURN_IF(!declared(node.domain), "Replacement domain has no rule opset requirement: ", node.domain); + } else { + ORT_RETURN_IF(!declared(recipe.domain), "Replacement domain has no rule opset requirement: ", recipe.domain); + } + ORT_RETURN_IF_ERROR(Replace(graph, *winning_pattern, winning_match, recipe, producers, added)); + ++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) worklist.push_back(index); + for (auto it = node->OutputEdgesBegin(); it != node->OutputEdgesEnd(); ++it) { + next.push_back(it->GetNode().Index()); + } + } + affected = std::move(next); + } + } + 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..26284b7f09010 --- /dev/null +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -0,0 +1,205 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { +namespace rewrite { + +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; + 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; +}; + +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: + Pattern() = default; + explicit Pattern(std::string_view text); + Value Input(std::string_view name = ""); + Value Constant(float value); + Value Constant(int64_t value); + Value Constant(int value) { return Constant(int64_t{value}); } + Value Constant(double value) { return Constant(static_cast(value)); } + Value NamedValue(std::string_view name) const; + void Name(Value value, std::string_view name); + void Attributes(std::initializer_list matchers); + void Attributes(InlinedVector matchers); + Value Op(Operator op, std::initializer_list inputs); + template , int> = 0> + Value Op(Name&& name, std::initializer_list inputs, std::string_view domain = "") { + return Op(Operator{std::string(std::string_view(name)), std::string(domain)}, inputs); + } + InlinedVector MultiOp(Operator op, std::initializer_list inputs, size_t output_count); + InlinedVector MultiOp(Operator op, gsl::span inputs, size_t output_count); + template , int> = 0> + InlinedVector MultiOp(Name&& name, std::initializer_list inputs, size_t output_count, + std::string_view domain = "") { + return MultiOp(Operator{std::string(std::string_view(name)), std::string(domain)}, inputs, output_count); + } + + 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_; } + + private: + friend struct Match; + friend class PatternRewriter; + size_t value_count_ = 0; + InlinedVector nodes_; + InlinedVector values_; + InlinedHashMap> constants_; + Value AddValue(std::string_view name); +}; + +struct Match { + InlinedVector nodes; + InlinedVector values; + const NodeArg* valueOf(std::string_view name) const; + void update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) const; + const ONNX_NAMESPACE::AttributeProto& attributeOf(std::string_view name) const; + void ClearUpdates() const { extra_attributes_.clear(); } + const Pattern* pattern = nullptr; + + private: + mutable NodeAttributes extra_attributes_; +}; + +struct Replacement { + struct Output { + size_t node; + size_t slot = 0; + }; + struct Scalar { + std::variant 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; + }; + Replacement() = default; + Replacement(std::string op_name, std::string op_domain, InlinedVector arguments, NodeAttributes attrs) + : op(std::move(op_name)), + domain(std::move(op_domain)), + inputs(std::move(arguments)), + attributes(std::move(attrs)) {} + explicit Replacement(std::string_view text); + explicit Replacement(InlinedVector nodes); + // 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_; } + std::string op; + std::string domain; + InlinedVector inputs; + NodeAttributes attributes; + + private: + friend class PatternRewriter; + std::shared_ptr> plan_; +}; + +struct Rule { + std::string name; + Pattern pattern; + InlinedVector opsets; + // Guards must be pure. Recipes may set extra attribute bindings, but must not mutate the graph + // or retain pointers into a Match. + std::function guard; + std::function replacement; +}; + +using Rules = InlinedVector; + +struct Statistics { + size_t visited = 0; + size_t candidates = 0; + size_t structural_checks = 0; + size_t guards = 0; + size_t matches = 0; + size_t rewrites = 0; +}; + +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}; +}; + +// 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_; +}; + +// Float32 examples adapted from genuine ONNXScript rules; EP eligibility is controlled by Options. +Rules MakeExampleRules(); + +} // namespace rewrite +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc new file mode 100644 index 0000000000000..9e8da4f42a1cf --- /dev/null +++ b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc @@ -0,0 +1,189 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include "core/optimizer/pattern_rewriter.h" + +#include +#include + +#include "core/graph/graph_utils.h" +#include "core/optimizer/initializer.h" + +namespace onnxruntime { +namespace rewrite { +namespace { +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; +} + +bool Float32(const Match& match) { + for (const NodeArg* value : match.values) { + const auto* type = value->TypeAsProto(); + if (!type || !type->has_tensor_type() || + type->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) + return false; + } + return true; +} + +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; +} +} // namespace + +Rules MakeExampleRules() { + Rules rules; + { + Pattern p(R"( + (x) => + add = Add(x, 3.0) + clip = Clip(add, 0.0, 6.0) + mul = Mul(clip, x) + return Div(mul, 6.0) + )"); + auto guard = [](const Graph&, const Match& m) { return Float32(m); }; + Replacement replacement("(x) => HardSwish(x)"); + Rule rule{"HardSwish", p, {{kOnnxDomain, 14, 21}}, guard, replacement}; + rules.push_back(std::move(rule)); + } + { + Pattern p(R"( + (x) => + sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) + return Mul(sigmoid, x) + )"); + auto guard = [](const Graph&, const Match& m) { return Float32(m); }; + Replacement replacement("(x) => HardSwish(x)"); + Rule rule{"HardSigmoidMul", p, {{kOnnxDomain, 14, 21}}, guard, replacement}; + rules.push_back(std::move(rule)); + } +#ifndef DISABLE_CONTRIB_OPS + { + Pattern p("(x, y, divisor) => mm = MatMul(x, y); return Div(mm, divisor)"); + auto guard = [](const Graph& graph, const Match& m) { + auto scalar = Scalar(graph, *m.valueOf("divisor")); + return Float32(m) && scalar && *scalar != 0.0f && std::isfinite(1.0f / *scalar); + }; + Replacement plan("(x, y, alpha) => com.microsoft.FusedMatMul(x, y, alpha=alpha)"); + auto replacement = [plan](const Graph& graph, const Match& m) { + m.update("alpha", FloatAttribute("alpha", 1.0f / *Scalar(graph, *m.valueOf("divisor")))); + return plan; + }; + Rule rule{"MatMulDiv", p, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}, guard, replacement}; + rules.push_back(std::move(rule)); + } + { + Pattern p("(x, y) => transpose = Transpose(y); return MatMul(x, transpose)"); + auto guard = [](const Graph&, const Match& m) { + if (!Float32(m)) 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; + }; + Replacement replacement("(x,y) => com.microsoft.FusedMatMul(x,y,transB=1)"); + Rule rule{"TransposeMatMul", p, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}, guard, replacement}; + rules.push_back(std::move(rule)); + } + { + Pattern p("(x, y, divisor) => mm = com.microsoft::FusedMatMul(x, y); return Div(mm, divisor)"); + auto guard = [](const Graph& graph, const Match& m) { + auto scalar = Scalar(graph, *m.valueOf("divisor")); + auto it = m.nodes[0]->GetAttributes().find("alpha"); + float alpha = it == m.nodes[0]->GetAttributes().end() ? 1.0f : it->second.f(); + return Float32(m) && scalar && *scalar != 0.0f && std::isfinite(alpha / *scalar); + }; + Replacement plan(R"( + (x, y, alpha, transA, transB, transBatchA, transBatchB) => + com.microsoft.FusedMatMul(x, y, alpha=alpha, transA=transA, transB=transB, + transBatchA=transBatchA, transBatchB=transBatchB) + )"); + auto replacement = [plan](const Graph& graph, const Match& m) { + 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(); + m.update("alpha", FloatAttribute("alpha", alpha / *Scalar(graph, *m.valueOf("divisor")))); + 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 plan; + }; + Rule rule{"FusedMatMulDiv", p, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}, guard, replacement}; + rules.push_back(std::move(rule)); + } + { + Pattern p(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 guard = [](const Graph& graph, const Match& m) { + if (!Float32(m)) 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"); + return 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; + }; + Replacement plan(R"( + (x, scale, fused_epsilon) => + SimplifiedLayerNormalization(x, scale, axis=-1, epsilon=fused_epsilon, stash_type=1) + )"); + auto replacement = [plan](const Graph& graph, const Match& m) { + m.update("fused_epsilon", FloatAttribute("epsilon", *Scalar(graph, *m.valueOf("epsilon")))); + return plan; + }; + Rule rule{"RmsNormalization", p, {{kOnnxDomain, 7, 17}}, guard, replacement}; + rules.push_back(std::move(rule)); + } +#endif + return rules; +} + +} // namespace rewrite +} // 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..82a2a2171a11d --- /dev/null +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -0,0 +1,1989 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#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; + +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) { + Pattern p("(x) => a = Identity(x); return Identity(a)"); + auto guard = [accept](const Graph&, const Match&) { return accept; }; + auto replacement = [](const Graph&, const Match& m) { + return Replacement{"Identity", kOnnxDomain, {m.valueOf("x")}, {}}; + }; + Rule rule{name, p, {{kOnnxDomain, 1, 21}}, guard, replacement}; + return rule; +} + +Rule IdentitySplit() { + Pattern p; + auto x = p.Input(); + auto first = p.Op("Identity", {x}); + p.MultiOp("Split", {first}, 2); + auto guard = [](const Graph&, const Match&) { return true; }; + auto replacement = [x](const Graph&, const Match& match) { + return Replacement{"Split", kOnnxDomain, {x}, match.nodes.back()->GetAttributes()}; + }; + Rule rule{"IdentitySplit", p, {{kOnnxDomain, 13, 17}}, guard, replacement}; + return rule; +} + +Rule SplitConcat() { + Pattern p("(x) => a, b = Split(x); return Concat(a, b)"); + auto guard = [](const Graph&, const 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; + }; + auto replacement = [](const Graph&, const Match& m) { + return Replacement{"Identity", kOnnxDomain, {m.valueOf("x")}, {}}; + }; + Rule rule{"SplitConcat", p, {{kOnnxDomain, 13, 17}}, guard, replacement}; + return rule; +} + +// Multi-output adaptation of ONNXScript's LayerNormBiasFusion; unused statistics may be removed. +Rule LayerNormBias() { + Pattern p("(x, scale, bias) => y, mean, inv_std = LayerNormalization(x, scale); return Add(y, bias)"); + auto guard = [](const Graph&, const 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"); + return 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); + }; + auto replacement = [](const Graph&, const Match& match) { + return Replacement{"LayerNormalization", + "", + {match.valueOf("x"), match.valueOf("scale"), match.valueOf("bias")}, + match.nodes[0]->GetAttributes()}; + }; + Rule rule{"LayerNormBias", p, {{kOnnxDomain, 17, 17}}, guard, replacement}; + return rule; +} + +// Small-weight eager decoder: B=1, Hq=4, Hkv=2, D=4. Cache length is 3. +// See the design document for the model/export sources and deliberate simplifications. +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, NamedValuesAndAttributePredicates) { + Pattern p; + auto x = p.Input("x"); + auto first = p.Op("HardSigmoid", {x}); + p.Name(first, "sigmoid"); + p.Attributes({AttributeMatcher::Equal("alpha", 0.25f), + {"beta", [](const auto& attr) { + return attr.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT && attr.f() > 0.4f; + }}}); + p.Op("Mul", {first, x}); + 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(p.Input("x"), OnnxRuntimeException); + EXPECT_THROW(p.NamedValue("missing"), OnnxRuntimeException); + EXPECT_THROW(p.Name(x, "sigmoid"), OnnxRuntimeException); + EXPECT_THROW(p.Op("Identity", {Value{-1}}), 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; + 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); return Identity(middle)"); + first.guard = [=](const Graph&, const 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.guard = [=](const Graph&, const 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, output](const Graph&, const Match& m) { + const NodeArg* input = scenario == 0 ? m.valueOf("input") + : scenario == 1 ? m.valueOf("middle") + : scenario == 2 ? output + : nullptr; + return Replacement{"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(); + Pattern p; + auto x = p.Input("x"); + auto a = p.Op("Identity", {x}); + auto b = p.Op("Identity", {a}); + p.MultiOp("Split", {b}, 2); + split.pattern = p; + split.replacement = Replacement("(x) => a=Identity(x); return Split(a, axis=0)"); + 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)"}) + 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_THROW(PatternRewriter({invalid}), OnnxRuntimeException); + 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 (int scenario = 0; scenario < 9; ++scenario) { + 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]; + InlinedVector nodes{{"Identity", "", {std::string("x")}, {}, 1}, + {"HardSwish", "", {Replacement::Output{0}}, {}, 1}}; + if (scenario == 0) nodes[1].op = "NoSuchReplacementOperator"; + if (scenario == 1) nodes[1].inputs = {Replacement::Output{1}}; + if (scenario == 2) nodes[1].inputs = {Replacement::Output{0, 1}}; + if (scenario == 3) nodes[1].inputs = {std::string("x")}; + if (scenario == 4) { + ONNX_NAMESPACE::AttributeProto attr; + attr.set_name("axis"); + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); + attr.set_s("wrong"); + nodes[1].op = "Softmax"; + nodes[1].attributes.emplace("axis", attr); + } + if (scenario == 5) nodes[1].domain = kMSDomain; + if (scenario == 6) nodes[0].output_count = 0; + if (scenario == 7) { + nodes[0].op = "Add"; + nodes[0].inputs.push_back(Replacement::Scalar{0.0f}); + nodes[1].op = "NoSuchReplacementOperator"; + } + rule.replacement = [recipe = Replacement(std::move(nodes))](const Graph&, const Match&) { return recipe; }; + if (scenario == 8) { + rule.replacement = [](const Graph&, const Match&) { + Replacement recipe("(x)=>HardSwish(x)"); + recipe.op = "Identity"; + 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, ReplacementScalarConstantsAndProgrammaticSequence) { + 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(InlinedVector{ + {"Identity", "", {std::string("x")}, {}, 1}, {"Identity", "", {Replacement::Output{0}}, {}, 1}}); + 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.guard = [](const Graph&, const Match&) { return true; }; + second.replacement = Replacement("(second)=>a=Identity(second);return Identity(a)"); + auto followup = IdentityPair(); + followup.guard = [](const Graph&, const 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)"); + }; + 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(test.graph.NumberOfNodes(), 4u); + ASSERT_STATUS_OK(test.graph.Resolve()); + } +} + +TEST(PatternRewriterTests, ReplacementExtraAttributeBindings) { + 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_THROW(Pattern{"(x,a)=>return HardSigmoid(x,alpha=a)"}, OnnxRuntimeException); + 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; + // Deliberate guard misuse verifies that rejected candidates cannot leak updates. + first.guard = [alpha](const Graph&, const Match& m) { + m.update("alpha", alpha); + return false; + }; + Replacement plan("(x,alpha,beta)=>h=HardSigmoid(x,alpha=alpha,beta=beta);return Mul(h,x)"); + second.replacement = [plan, scenario, make_attribute](const Graph&, const Match& m) { + auto beta = make_attribute(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + beta.set_f(0.5f); + m.update("beta", beta); + if (scenario == 1) return plan; // 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); + if (scenario == 4) { + auto malformed = plan; + malformed.attributes.emplace("extra", beta); + return malformed; + } + return plan; + }; + 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)"); + Replacement plan("(x,axis)=>Softmax(x,axis=axis)"); + int calls = 0; + rule.replacement = [plan, &calls](const Graph&, const Match& m) { + ONNX_NAMESPACE::AttributeProto axis; + axis.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + axis.set_i(calls++ == 0 ? 0 : 1); + m.update("axis", axis); + return plan; + }; + 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, repeat); + EXPECT_EQ(test.graph.NumberOfNodes(), 3u); + EXPECT_EQ(stats.matches, 1u); + ASSERT_STATUS_OK(rewriter.Run(test.graph, stats)); + 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)"); + Replacement plan("(x,perm)=>Transpose(x,perm=perm)"); + auto replacement = [plan](const Graph&, const 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 plan; + }; + rule.replacement = replacement; + 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, 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].guard](const Graph& graph, const 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.guard = 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.guard = 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; + auto integer = p.Constant(3); + auto floating = p.Constant(3.0); + EXPECT_TRUE(std::holds_alternative(p.Constants().at(integer))); + EXPECT_TRUE(std::holds_alternative(p.Constants().at(floating))); +} + +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.guard = [&guards](const Graph&, const 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.guard = [&first](const Graph&, const Match&) { + ++first; + return true; + }; + c.guard = [&second](const Graph&, const 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, InvalidRuleAndRecipe) { + Pattern disconnected; + auto x = disconnected.Input(); + disconnected.Op({"Identity", ""}, {x}); + disconnected.Op({"Identity", ""}, {x}); + auto rule = IdentityPair(); + rule.pattern = disconnected; + EXPECT_THROW(PatternRewriter({rule}), 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 (int scenario = 0; scenario < 5; ++scenario) { + rule.replacement = [scenario](const Graph&, const Match&) { + Replacement recipe{"Identity", "", {0}, {}}; + if (scenario == 0) recipe.op = "NoSuchReplacementOperator"; + if (scenario == 1) recipe.domain = "not.imported"; + if (scenario == 2) recipe.inputs = {1}; // Deleted intermediate, not an external input. + if (scenario == 3) { + ONNX_NAMESPACE::AttributeProto attr; + attr.set_name("invalid"); + attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + attr.set_i(1); + recipe.attributes.emplace("invalid", attr); + } + if (scenario == 4) recipe.domain = kMSDomain; // Imported, but not declared by this rule. + 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}}; + Pattern pattern; + auto input = pattern.Input(); + auto first = pattern.Op({"Identity", kOnnxDomainAlias}, {input}); + pattern.Op({"Identity", kOnnxDomainAlias}, {first}); + rule.pattern = std::move(pattern); + int guards = 0; + rule.guard = [&guards](const Graph&, const Match&) { + ++guards; + return true; + }; + rule.replacement = [](const Graph&, const Match&) { return Replacement{"Identity", kOnnxDomainAlias, {0}, {}}; }; + // 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.guard = [&inactive_guards](const Graph&, const Match&) { + ++inactive_guards; + return true; + }; + active.guard = [&active_guards](const Graph&, const Match&) { + ++active_guards; + return true; + }; + later.guard = [&later_guards](const Graph&, const 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.guard = [&guards](const Graph&, const 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, 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.guard = [](const Graph& graph, const Match& m) { + return graph.NodeProducesGraphOutput(*m.nodes.back()); + }; + Pattern p; + auto input = p.Input(); + auto first = p.Op({"Identity", ""}, {input}); + auto second = p.Op({"Identity", ""}, {first}); + p.Op({"Identity", ""}, {second}); + auto guard = [](const Graph&, const Match&) { return true; }; + auto replacement = [input](const Graph&, const Match&) { return Replacement{"Identity", "", {input}, {}}; }; + Rule higher_priority{"higher_priority", p, {{kOnnxDomain, 1, 21}}, guard, replacement}; + 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; + auto input = p.Input(); + auto identity = p.Op({"Identity", ""}, {input}); + p.Op({"Dropout", ""}, {identity}); + auto rule = IdentityPair(); + rule.pattern = p; + Statistics stats; + ASSERT_STATUS_OK(PatternRewriter({rule}).Run(test.graph, stats)); + EXPECT_EQ(stats.rewrites, 0u); +} + +TEST(PatternRewriterTests, MultiOpDefinitionsAndOverloads) { + Pattern p; + auto x = p.Input(); + std::string identity = "Identity"; + auto a = p.Op(identity, {x}); + auto b = p.Op(std::string_view(identity), {a}, kOnnxDomainAlias); + auto outputs = p.MultiOp({"Split", kOnnxDomainAlias}, {b}, 2); + 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.MultiOp("Split", {b}, 0), OnnxRuntimeException); + EXPECT_THROW(p.MultiOp("Split", {b}, 65), OnnxRuntimeException); + EXPECT_THROW(p.MultiOp("Split", {999}, 2), OnnxRuntimeException); + EXPECT_EQ(p.ValueCount(), count); + for (int i = 0; i < 3; ++i) p.MultiOp("Split", {b}, 64); + EXPECT_THROW(p.MultiOp("Split", {b}, 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; + auto input = pattern.Input(); + auto outputs = pattern.MultiOp("Split", {input}, 2); + pattern.Op("Identity", {outputs[1]}); + auto rule = SplitConcat(); + rule.pattern = pattern; + rule.guard = [](const Graph&, const 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()); + auto invalid = IdentitySplit(); + invalid.replacement = [](const Graph&, const Match&) { return Replacement{"Identity", "", {0}, {}}; }; + Statistics stats; + 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"); + Pattern single_output; + auto input = single_output.Input(); + auto first = single_output.Op("Identity", {input}); + single_output.Op("Split", {first}); + wrong_count.pattern = single_output; + ASSERT_STATUS_OK(PatternRewriter({wrong_count, IdentitySplit()}, {.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, 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()); + Pattern p; + auto px = p.Input(), py = p.Input(); + auto psum = p.Op({"Add", ""}, {px, py}); + auto pproduct = p.Op({"Mul", ""}, {psum, px}); + p.Op({"Add", ""}, {psum, pproduct}); + auto guard = [](const Graph&, const Match&) { return true; }; + auto replacement = [](const Graph&, const Match&) { return Replacement{"Identity", "", {0}, {}}; }; + Rule rule{"diamond", p, {{kOnnxDomain, 7, 21}}, guard, replacement}; + 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, 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]; + 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 + +} // namespace test +} // namespace onnxruntime From 07afc0666b1faa1f1a88ec51b885bcc12802be07 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 13 Oct 2026 18:01:16 -0700 Subject: [PATCH 02/45] Check example rewrite types on relevant value bindings Replace match-wide float32 scans with explicit data-binding checks. Keep scalar initializer validation and add focused guard tests plus numerical RMS coverage with an independently typed int64 Pow exponent. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 32 ++++-- .../core/optimizer/pattern_rewriter_rules.cc | 26 +++-- .../test/optimizer/pattern_rewriter_test.cc | 102 ++++++++++++++++++ 3 files changed, 137 insertions(+), 23 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 5fa61a9fdcc2f..f980d7fb0017e 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -4,7 +4,8 @@ It is **not registered in the default optimization pipeline**, and introduces no session configuration switch or public C ABI. C++ callers may register it using `InferenceSession::RegisterGraphTransformer`, just like existing custom transformers. -The example predicates are provider-independent and restrict types to float32. +The example predicates are provider-independent and restrict relevant data bindings +to float32, rather than scanning every matched value. The engine's generic partition policy defaults to CPU or unassigned graphs; allowing other EPs requires the caller to verify replacement kernel compatibility. @@ -52,13 +53,9 @@ p.Attributes({AttributeMatcher::Equal("alpha", 1.0f / 6.0f), AttributeMatcher::Equal("beta", 0.5f)}); p.Op("Mul", {h, x}); auto guard = [](const Graph&, const Match& m) { - for (const NodeArg* value : m.values) { - const auto* type = value->TypeAsProto(); - if (!type || !type->has_tensor_type() || - type->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) - return false; - } - return true; + const auto* type = m.valueOf("x")->TypeAsProto(); + return type && type->has_tensor_type() && + type->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT; }; Replacement replacement("(x) => HardSwish(x)"); Rule rule{ @@ -167,11 +164,28 @@ Pattern p(R"( )"); ``` -Its guard only checks float32, since the pattern checks every fixed scalar. +Its guard only checks that `x` is float32, since the pattern checks every fixed scalar. RMSNorm similarly uses `Pow(x, 2.0)` directly; epsilon remains a named binding because its value varies and is copied to the fused operator's attributes. MatMul division factors also remain variable for computed scaling attributes. +The six example guards check types on these specific bindings: + +| Rule | Float32 bindings | +|---|---| +| HardSwish | `x` | +| HardSigmoidMul | `x` | +| MatMulDiv | `x`, `y`; `divisor` is checked as a float32 scalar initializer | +| TransposeMatMul | `x`, `y` | +| FusedMatMulDiv | `x`, `y`; `divisor` is checked as a float32 scalar initializer | +| RmsNormalization | `x`, `scale`; `epsilon` is checked as a float32 scalar initializer | + +Resolved operator schemas establish the corresponding intermediate/output types. +Guards do not impose a match-wide element type: independently typed operands such +as a Pow exponent or shape/index values need their own semantic constraints. +The supplied RMSNorm pattern still requires the float32 literal `2.0` exponent; +that literal constraint is separate from the guard's data-type checks. + Attributes use Python keyword syntax, after positional inputs: `h = HardSigmoid(x, alpha = 0.16666667, beta = 0.5)`. Signed/scientific numbers and single/double-quoted string attributes are supported. diff --git a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc index 9e8da4f42a1cf..bdf691a325a37 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc @@ -19,14 +19,10 @@ ONNX_NAMESPACE::AttributeProto FloatAttribute(const char* name, float value) { return attr; } -bool Float32(const Match& match) { - for (const NodeArg* value : match.values) { - const auto* type = value->TypeAsProto(); - if (!type || !type->has_tensor_type() || - type->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) - return false; - } - return true; +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; } std::optional Scalar(const Graph& graph, const NodeArg& value) { @@ -58,7 +54,7 @@ Rules MakeExampleRules() { mul = Mul(clip, x) return Div(mul, 6.0) )"); - auto guard = [](const Graph&, const Match& m) { return Float32(m); }; + auto guard = [](const Graph&, const Match& m) { return IsFloat32(*m.valueOf("x")); }; Replacement replacement("(x) => HardSwish(x)"); Rule rule{"HardSwish", p, {{kOnnxDomain, 14, 21}}, guard, replacement}; rules.push_back(std::move(rule)); @@ -69,7 +65,7 @@ Rules MakeExampleRules() { sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) return Mul(sigmoid, x) )"); - auto guard = [](const Graph&, const Match& m) { return Float32(m); }; + auto guard = [](const Graph&, const Match& m) { return IsFloat32(*m.valueOf("x")); }; Replacement replacement("(x) => HardSwish(x)"); Rule rule{"HardSigmoidMul", p, {{kOnnxDomain, 14, 21}}, guard, replacement}; rules.push_back(std::move(rule)); @@ -79,7 +75,8 @@ Rules MakeExampleRules() { Pattern p("(x, y, divisor) => mm = MatMul(x, y); return Div(mm, divisor)"); auto guard = [](const Graph& graph, const Match& m) { auto scalar = Scalar(graph, *m.valueOf("divisor")); - return Float32(m) && scalar && *scalar != 0.0f && std::isfinite(1.0f / *scalar); + return IsFloat32(*m.valueOf("x")) && IsFloat32(*m.valueOf("y")) && + scalar && *scalar != 0.0f && std::isfinite(1.0f / *scalar); }; Replacement plan("(x, y, alpha) => com.microsoft.FusedMatMul(x, y, alpha=alpha)"); auto replacement = [plan](const Graph& graph, const Match& m) { @@ -92,7 +89,7 @@ Rules MakeExampleRules() { { Pattern p("(x, y) => transpose = Transpose(y); return MatMul(x, transpose)"); auto guard = [](const Graph&, const Match& m) { - if (!Float32(m)) return false; + 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"); @@ -118,7 +115,8 @@ Rules MakeExampleRules() { auto scalar = Scalar(graph, *m.valueOf("divisor")); auto it = m.nodes[0]->GetAttributes().find("alpha"); float alpha = it == m.nodes[0]->GetAttributes().end() ? 1.0f : it->second.f(); - return Float32(m) && scalar && *scalar != 0.0f && std::isfinite(alpha / *scalar); + return IsFloat32(*m.valueOf("x")) && IsFloat32(*m.valueOf("y")) && + scalar && *scalar != 0.0f && std::isfinite(alpha / *scalar); }; Replacement plan(R"( (x, y, alpha, transA, transB, transBatchA, transBatchB) => @@ -153,7 +151,7 @@ Rules MakeExampleRules() { return Mul(normalize, scale) )"); auto guard = [](const Graph& graph, const Match& m) { - if (!Float32(m)) return false; + 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(); diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 82a2a2171a11d..6cd6b34d84aad 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1744,6 +1744,77 @@ TEST(PatternRewriterTests, DoubleTypeRejection) { 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); + TestGraph test; + auto& b = test.builder; + auto* non_float = b.MakeScalarInitializer(2.0); + auto* integer = b.MakeScalarInitializer(2); + auto* unknown = b.MakeIntermediate(); + 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 guard = rule.guard; + size_t checked = 0; + rule.guard = [&](const Graph& graph, const Match& match) { + EXPECT_TRUE(guard(graph, match)); + for (auto binding : bindings) { + SCOPED_TRACE(binding); + // Exercise the guard 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(guard(graph, changed)); + } + } + ++checked; + return guard(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; @@ -1813,6 +1884,37 @@ TEST(PatternRewriterTests, SubgraphsAndImplicitCaptureBoundary) { } #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; From a62df40f7edad31db003e0d0430673a5ad837577 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 13 Oct 2026 18:15:04 -0700 Subject: [PATCH 03/45] Make rewrite definitions parser-only and document matcher semantics Expose textual Pattern and Replacement construction with read-only introspection. Keep parser builders private through a named internal parser, support explicit named final outputs, migrate correctness helpers and malformed-recipe coverage to parsed definitions, and replace the design investigation with a user-facing rule-writing guide. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 853 +++++++----------- .../core/optimizer/pattern_rewriter.cc | 58 +- onnxruntime/core/optimizer/pattern_rewriter.h | 51 +- .../test/optimizer/pattern_rewriter_test.cc | 274 +++--- 4 files changed, 463 insertions(+), 773 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index f980d7fb0017e..2b6568f242fd3 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -1,607 +1,376 @@ -# Experimental indexed C++ pattern rewriter +# Writing C++ graph rewrite rules -`onnxruntime/core/optimizer/pattern_rewriter.h` defines an opt-in `GraphTransformer`. -It is **not registered in the default optimization pipeline**, and introduces no -session configuration switch or public C ABI. C++ callers may register it using -`InferenceSession::RegisterGraphTransformer`, just like existing custom transformers. -The example predicates are provider-independent and restrict relevant data bindings -to float32, rather than scanning every matched value. -The engine's generic partition policy defaults to CPU or unassigned graphs; -allowing other EPs requires the caller to verify replacement kernel compatibility. +The experimental `onnxruntime::rewrite::PatternRewriter` replaces connected +operator patterns with smaller, 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. +Guards and replacement callbacks remain ordinary C++ functions. -## Source investigation and supported examples +A `Rule` combines a diagnostic name, a parsed pattern, domain/opset requirements, +a semantic guard, and a parsed replacement or callback returning one. The matcher +establishes structural and deletion safety. **The rule author must establish +semantic equivalence, supported types, and kernel compatibility.** -The ONNXScript reference revision is -[`9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36`](https://github.com/microsoft/onnxscript/tree/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36). -The following files were examined, rather than assuming that Python syntax defines -the right architecture for ORT: +## Example: HardSigmoid followed by Mul -| ONNXScript source at that revision | C++ adaptation | -|---|---| -| [`rewriter/_matcher.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/_matcher.py) | Backward DAG binding, repeated-value equality, distinct matched nodes, and removal boundary checks | -| [`rewriter/_rewrite_rule.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/_rewrite_rule.py) | Separate structural patterns, semantic guards, replacement recipes, ordered priority, recursive subgraph handling | -| [`rules/common/_fuse_hardswish.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/rules/common/_fuse_hardswish.py) | `Div(Mul(Clip(Add(x,3),0,6),x),6)` and `Mul(HardSigmoid(x),x)` to HardSwish | -| [`ort_fusions/fused_matmul_rule_sets.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/ort_fusions/fused_matmul_rule_sets.py) | MatMul/scalar Div, FusedMatMul/scalar Div (preserving attributes), and last-two-dimension Transpose/MatMul | -| [`ort_fusions/rms_normalization.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/ort_fusions/rms_normalization.py) | Float32 Pow/ReduceMean/Add/Sqrt/Reciprocal/Mul/Mul to SimplifiedLayerNormalization | -| [`ort_fusions/softmax.py`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/ort_fusions/softmax.py) | Studied but deliberately not enabled: fp16 cast removal needs a separately verified precision/EP policy | -| [`rules/fusion/_layer_norm.py`, `LayerNormBiasFusion`](https://github.com/microsoft/onnxscript/blob/9ff5a7c7718ab99d2ee0fa9b5effee2c543f3e36/onnxscript/rewriter/rules/fusion/_layer_norm.py) | Test-only adaptation matches all three LayerNormalization outputs, fuses a trailing bias Add, and rejects observable statistics | - -Existing ORT `SelectorActionRegistry` already dispatches selectors by root op/domain -and version. `QuickGeluFusion` illustrates repeated inputs and provider restrictions, -while `graph_utils::FinalizeNodeFusion` illustrates edge maintenance. The new -transformer uses ORT's graph/node/edge APIs directly; it does not copy the graph to -a parallel IR. Generic backward bindings and sharing identical structural plans -are the additional capabilities over hand-written selectors. - -## Defining a rule - -The small DSL uses `Value {std::string name; int handle;}`. A name helps diagnose -patterns; only the handle determines binding identity. An input introduces a variable; -`Op` introduces one output, while `MultiOp` introduces an ordered output list. -Reusing a handle requires the identical -`NodeArg`, not merely the same shape or equivalent contents. Pattern operators -specify only name and domain. **Model-imported opset ranges belong to the rule**, -not to individual nodes or their schema since-versions. +This complete rule recognizes the HardSwish decomposition with the exact +HardSigmoid attributes shown: ```cpp +#include "core/optimizer/pattern_rewriter.h" + +using namespace onnxruntime; using namespace onnxruntime::rewrite; -Pattern p; -auto x = p.Input("x"); -auto h = p.Op("HardSigmoid", {x}); -p.Name(h, "sigmoid"); -p.Attributes({AttributeMatcher::Equal("alpha", 1.0f / 6.0f), - AttributeMatcher::Equal("beta", 0.5f)}); -p.Op("Mul", {h, x}); -auto guard = [](const Graph&, const Match& m) { - const auto* type = m.valueOf("x")->TypeAsProto(); + +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; -}; -Replacement replacement("(x) => HardSwish(x)"); -Rule rule{ - "HardSigmoidMul", p, - {{kOnnxDomain, 14, 21}}, // inclusive imported opset range - guard, replacement}; -auto transformer = std::make_unique(Rules{rule}); -ORT_RETURN_IF_ERROR(session.RegisterGraphTransformer(std::move(transformer), - TransformerLevel::Level1)); -``` - -`pattern_rewriter_rules.cc` contains complete readable guard/recipe examples. -Each example defines its pattern, guard and replacement separately, assembles a -named `Rule rule`, and registers it with `rules.push_back(std::move(rule))`. -`Op(name, inputs)` defaults to the standard ONNX domain; the optional third -argument selects another domain, e.g. `p.Op("FusedMatMul", {x, y}, kMSDomain)`. -The original `Op(Operator{name, domain}, inputs)` overload remains available. -Name overloads accept literals, `std::string` and `std::string_view`; template -deduction keeps the original braced operator form unambiguous. - -`p.NamedValue("sigmoid")` retrieves a named handle. Nonempty names must be unique. -`p.Name(value, name)` updates the canonical name and existing node definitions; -subsequent operators also use that canonical name. Previously returned `Value` -copies are not renamed in place. `Value` converts to `int` for indexing -`Match::values`; an unnamed value has an empty name, and a default value has -invalid handle -1. - -In guards and replacement recipes, `m.valueOf("x")` directly returns the matched -`const NodeArg*`; there is no need to capture `p.NamedValue("x")` and index -`m.values` separately. Lookup also supports named intermediate outputs. -Unknown/empty names or unbound values throw an explicit error rather than -returning null. Names come from the current rule, even when its structural plan -is shared with a differently named pattern. Lookup scans the bounded pattern -value table and does not allocate a name map for each candidate. - -`Replacement::inputs` accepts either `Value` handles or matched `const NodeArg*` -bindings, so recipes can use `{m.valueOf("x"), m.valueOf("y")}` directly. -Pointer inputs are validated against the match's external input bindings: -intermediates, root outputs, nulls and unrelated graph values are rejected before -mutation. Matches reference immutable compiled patterns; do not retain them or -their pointers beyond a callback. - -### Attribute matchers and literal constants - -`p.Attributes(...)` attaches matchers to the most recently defined node. -Each `AttributeMatcher` contains an attribute name and a boolean predicate on -`const AttributeProto&`. Missing attributes fail the match; schema defaults are -not synthesized. Equality factories check both attribute type and value: -`Equal("axis", 1)` requires INT, `Equal("alpha", 0.5f)` requires FLOAT, -and `Equal("mode", std::string("linear"))` requires STRING. Extra unmatched -attributes are permitted. Arbitrary predicates are also supported: - -```cpp -p.Attributes({{"axis", [](const ONNX_NAMESPACE::AttributeProto& value) { - return value.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_INT && - value.i() < 0; -}}}); +} + +Rule MakeHardSwishRule() { + Pattern pattern(R"( + (x) => + sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) + return Mul(sigmoid, x) + )"); + auto guard = [](const Graph&, const Match& match) { + return IsFloat32(*match.valueOf("x")); + }; + Replacement replacement("(x) => HardSwish(x)"); + return Rule{"HardSigmoidMul", std::move(pattern), + {{kOnnxDomain, 14, 21}}, guard, replacement}; +} ``` -Predicates are trusted, pure C++ callbacks, with the same locality and lifetime -requirements as rule guards. Duplicate/empty attribute names, empty callbacks, -or attachment before any node are errors. +Both occurrences of `x` must bind the **same `NodeArg`**, not just tensors with +equal shapes or contents. The guard checks the relevant data binding `x`; +resolved schemas establish the corresponding intermediate and output types. +It does not require every matched value to have the same element type. -`p.Constant(3.0f)` introduces a float32 scalar initializer constraint; -`p.Constant(int64_t{3})` introduces an int64 constraint. Integer literals in C++ -select int64; `double` arguments are converted to float32. Float constants must -be finite. Matching requires an immutable rank-zero initializer of the exact -type and value; a one-element vector, an overridable initializer, or external-data -initializer does not match. These leaf handles can be used as replacement inputs, -but matching never creates tensors. +These opset bounds are a conservative supported range for this example, not an +automatic statement about every schema version. Floating-point rewrites can +change evaluation order; validate equivalence within the intended tolerance. -### Textual patterns +## Pattern and replacement syntax -The explicit `Pattern(std::string_view text)` constructor parses a textual pattern. -The default constructor remains available for the programmatic builder. -Malformed text throws `OnnxRuntimeException` with a byte-offset diagnostic -(or aborts in exception-disabled builds, following ORT's throwing macros). -The requested branched pattern can be written directly: +Definitions start with named inputs and `=>`. Assignments create named operator +outputs; `return` identifies the final operator, called the root. Whitespace +separates calls, and semicolons are optional separators. Comments are not supported +inside definitions. ```cpp -Pattern p(R"( - (x, low, high, divisor) => +Pattern pattern(R"( + (x) => add = Add(x, 3.0) - clip = Clip(add, low, high) - mul = Mul(clip, x) - return Div(mul, divisor) + clip = Clip(add, 0.0, 6.0) + product = Mul(clip, x) + return result = Div(product, 6.0) )"); +Replacement replacement("(x) => HardSwish(x)"); ``` -The header declares named inputs. Assignment names label operator outputs. -Reusing `x` enforces exact binding identity. Decimal/scientific literals become -float32 scalar constants, while integer spellings become int64 constants. -This text describes structure, **not a semantically complete HardSwish rule**: -the rule guard must additionally establish float32, low=0, high=6 and divisor=6. -The numerical test retains those checks before executing the replacement. +The example has one root output, named `result`. 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. -The actual example rule eliminates those fixed bindings entirely: +Use `domain::Operator(...)` for qualified operators: ```cpp -Pattern p(R"( - (x) => - add = Add(x, 3.0) - clip = Clip(add, 0.0, 6.0) - mul = Mul(clip, x) - return Div(mul, 6.0) -)"); +Pattern pattern("(x,y) => mm = com.microsoft::FusedMatMul(x,y); return Identity(mm)"); +Replacement replacement("(x,y) => com.microsoft::FusedMatMul(x,y)"); ``` -Its guard only checks that `x` is float32, since the pattern checks every fixed scalar. -RMSNorm similarly uses `Pow(x, 2.0)` directly; epsilon remains a named binding -because its value varies and is copied to the fused operator's attributes. -MatMul division factors also remain variable for computed scaling attributes. +An unqualified operator uses the standard ONNX domain. `ai.onnx::Identity` +is equivalent to `Identity`. Replacements also accept +`com.microsoft.FusedMatMul(...)`; patterns use the `::` form. -The six example guards check types on these specific bindings: +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`, `ValueCount`, and `IsInput` +are available for introspection. -| Rule | Float32 bindings | -|---|---| -| HardSwish | `x` | -| HardSigmoidMul | `x` | -| MatMulDiv | `x`, `y`; `divisor` is checked as a float32 scalar initializer | -| TransposeMatMul | `x`, `y` | -| FusedMatMulDiv | `x`, `y`; `divisor` is checked as a float32 scalar initializer | -| RmsNormalization | `x`, `scale`; `epsilon` is checked as a float32 scalar initializer | - -Resolved operator schemas establish the corresponding intermediate/output types. -Guards do not impose a match-wide element type: independently typed operands such -as a Pow exponent or shape/index values need their own semantic constraints. -The supplied RMSNorm pattern still requires the float32 literal `2.0` exponent; -that literal constraint is separate from the guard's data-type checks. - -Attributes use Python keyword syntax, after positional inputs: -`h = HardSigmoid(x, alpha = 0.16666667, beta = 0.5)`. -Signed/scientific numbers and single/double-quoted string attributes are supported. -Only escaped quotes and backslashes are supported in strings. -The parsed spelling determines the attribute type: `axis=1` differs from `axis=1.0`. -An empty quoted string is valid. - -Multiple-output assignments and qualified domains are supported: - -```text -(x) => a, b = Split(x, axis=0); return Concat(a, b, axis=0) -(x, y) => mm = com.microsoft::FusedMatMul(x, y, alpha=0.5); return Identity(mm) +### 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. + +Keyword arguments constrain attributes, not tensor inputs: + +```cpp +Pattern pattern("(x) => h = HardSigmoid(x, alpha=0.25, beta=0.5); return Mul(h,x)"); ``` -Unqualified operators use the standard ONNX domain. Whitespace/newlines separate -tokens; semicolons may separate assignments. Values are ASCII identifiers. -Definitions must precede use; every input and node must reach the final operator. -There must be a final `return Op(...)`, with no subsequent statements. -All six example rules use the string constructor and retrieve bindings with -`Match::valueOf` for their guards and dynamic replacement recipes; static -replacement plans bind inputs once during compilation. Tests use the constructor for parsed -patterns and test-only Identity, Split/Concat and LayerNorm-bias rules; builder -tests and multi-output-root definitions retain the programmatic API. +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. -The parser checks duplicate names/attributes, undefined references, numeric range, -disconnected definitions and the same 64-node/256-value/64-input-output limits as -the C++ builder. Errors include a byte offset; `p = Pattern(text)` leaves `p` -unchanged if construction fails. -A one-node parsed pattern is valid, but cannot form a node-count-decreasing rule. +Pattern syntax supports scalar integer, float, and string attribute equality, +not attribute lists or custom predicate syntax. Inspect `match.nodes` in the C++ +guard for ranges, list-valued attributes, shapes, or other semantic conditions. +Positional inputs must precede keyword attributes. -The language intentionally excludes nested calls, lists/tensor literals, optional -placeholders, booleans, comments, `return` of an existing value and multi-output -`return` calls (use the C++ API for multi-output roots). It describes only matching: -opset requirements and semantic guards remain explicit C++. +Numeric literals in replacements synthesize rank-zero float32 or int64 +initializers. Replacement keyword literals set attributes on the new node. -### Reusable replacement plans +### Branched patterns and multiple outputs -`Replacement(std::string_view)` parses a reusable replacement definition, once: +Assignments can describe a DAG rather than only a chain: ```cpp -Replacement replacement("(x) => HardSwish(x)"); -Rule rule{"HardSwish", pattern, {{kOnnxDomain, 14, 21}}, guard, replacement}; +Pattern pattern(R"( + (x,y) => + sum = Add(x,y) + product = Mul(sum,x) + return Add(sum,product) +)"); ``` -The final operator call may omit `return`. Multi-node definitions use the same -assignment, domain and keyword-attribute syntax as patterns: +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 -Replacement replacement(R"( - (x) => - h = HardSigmoid(x, alpha=0.16666667, beta=0.5) - return Mul(h, x) -)"); +Pattern pattern("(x) => first,second = Split(x,axis=0); return Concat(first,second,axis=0)"); +Replacement replacement("(x) => Identity(x)"); ``` -This defines a two-node replacement, not a callback that parses a string on every -match. The constructor builds an immutable plan containing operator/domain names, -literal attributes, external input names and indexed earlier-output references. -Copies share that plan. When the plan is supplied directly as `Rule::replacement`, -transformer construction binds its external names to the rule's input handles, -and checks declared domains and node-count reduction. Binding is per rule, so -reusing a plan across differently ordered input headers is safe. - -At runtime the callback returns a shared plan reference. There is no parsing, -attribute-literal construction, or name lookup for these prebound inputs. -The engine binds the external `NodeArg`s, allocates fresh intermediate values, -creates nodes in definition order, and applies the stored attributes. -Callbacks that need computed attributes can use extra bindings with a captured -parsed plan (below), or return programmatic recipes. Do not construct -`Replacement(text)` inside a per-match callback. - -Replacement headers may select a subset of the matched pattern's inputs. -Tensor parameters cannot refer to matched intermediates. All declared replacement -parameters must be used as tensor inputs or attribute bindings, and all nodes -must reach the final operator. A multi-output assignment such as -`a, b = Split(x, axis=0)` declares the exact internal arity. -The final operator inherits **all matched root outputs in order**, including a -multi-output root; its schema must accept that count. Intermediate outputs cannot -be mapped independently to observable root slots yet. - -Numeric positional literals are synthesized as immutable rank-zero float32 or -int64 initializers, using the same literal typing as patterns: -`Replacement("(x) => Add(x, 0.0)")`. No Constant kernel or subsequent -constant-folding pass is required. There is no general tensor/list synthesis. - -Programmatic sequences use `NodeRecipe` and `Output{node_index, output_slot}`: +For a multi-output root, put the names on the final return assignment: ```cpp -Replacement replacement(InlinedVector{ - {"HardSigmoid", kOnnxDomain, {std::string("x")}, sigmoid_attributes, 1}, - {"Mul", kOnnxDomain, {Replacement::Output{0, 0}, std::string("x")}, {}, 1}}); +Pattern pattern("(x) => copied = Identity(x); return left,right = Split(copied,axis=0)"); +Replacement replacement("(x) => return left,right = Split(x,axis=0)"); ``` -`NodeRecipe` contains op, domain, inputs, attributes and internal output count. -It also has `attribute_bindings`, a map from target attribute names to extra -binding names. -Its inputs may be names, existing `Value` handles, matched `NodeArg*` bindings, -earlier-output references or `Replacement::Scalar{...}` literals. -The final node's output count is inherited from the root rather than this field. -The original four-argument single-node constructor remains supported, including -computed-attribute callbacks. Its public single-node fields cannot be combined -with an immutable parsed/programmatic plan; mixed representations are rejected. +This replacement preserves the root's two original output `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 only one output and cannot replace a two-output root. +Returning an existing value or a subset/reordering of an earlier operator's +outputs is not supported; the return must define the final operator. -#### Computed attributes through extra bindings +Intermediate operators may have unused outputs, but **none** of their outputs +may be exposed outside the matched deletion set. -Replacement keyword values may refer to names declared in its header. -Tensor arguments come from the pattern match; extra attribute arguments come -from `Match::update`, which accepts a concrete `AttributeProto`: +## Computed attributes with a replacement callback -```cpp -Replacement plan("(x, y, alpha) => com.microsoft.FusedMatMul(x, y, alpha=alpha)"); -auto replacement = [plan](const Graph& graph, const Match& m) { - m.update("alpha", FloatAttribute("alpha", - 1.0f / *Scalar(graph, *m.valueOf("divisor")))); - return plan; -}; -Rule rule{"MatMulDiv", pattern, opsets, guard, replacement}; -``` +Reuse a parsed replacement plan when an attribute depends on each match. +Declare the attribute binding in the replacement's input list, set it with +`match.update`, and reference it as a keyword argument. -Here `FloatAttribute` and `Scalar` are the helpers used by the example rules; -the guard establishes float32 and a finite, nonzero scalar divisor. -The plan is parsed **once when defining the rule**, not once per successful -match. Capturing/returning it shares the immutable plan. Computed attributes -are looked up and copied into the new nodes at rewrite time. Captured callback -plans still resolve tensor names at runtime; only directly supplied plans have -their tensor names prebound by transformer construction. - -`com.microsoft.FusedMatMul` is accepted in replacement text, as is the existing -`com.microsoft::FusedMatMul` spelling. Literal attributes and bound attributes -can be mixed, e.g. `axis=-1, epsilon=fused_epsilon, stash_type=1`. -Multiple nodes can share a binding, and target attribute names can differ from -binding names. Attribute proto names are normalized to the target names. -The proto carries the type, including list/tensor/string types not representable -as text literals; imported operator schemas validate the resolved attributes. - -Extra bindings are invocation-local state, separate from immutable matched -graph values. `update` is callable on the callback's `const Match&`, because it -changes only this extra state, not the graph or matched values. Repeated updates -to an extra name replace its value. Guards must remain pure; the engine clears -extra state between candidates and before invoking the winning recipe. No binding -is retained between rewrites, graph runs or rules; match-only mode never invokes -recipes. `attributeOf(name)` retrieves a previously set extra attribute. - -Empty names, undefined attribute types, reference attributes, missing bindings -and attempts to shadow any named pattern value are explicit errors. -A header name cannot be used both as a tensor argument and an attribute binding; -attribute references must name header parameters, not intermediate outputs. -All required extras are checked before graph mutation. Schema validation failures -use the same staged-node rollback as literal attributes. As with other recipes, -callbacks must establish semantic/type compatibility; extra bindings are not -an automatic semantic guard or a general graph transaction. - -The MatMulDiv and RMSNorm examples use captured plans with computed bindings. -TransposeMatMul uses a literal parsed replacement. FusedMatMulDiv also captures -a parsed plan, computing `alpha` and binding `transA`, `transB`, `transBatchA` -and `transBatchB` from the original node. Missing transpose attributes use the -imported schema's defaults; the replacement emits all five version-1 attributes -explicitly, preserving their effective values. All six examples therefore use -parsed replacements. Extra bindings describe individual attributes, not an -attribute-map spread; a rule for a schema with other attributes must explicitly -preserve those too. - -Every replacement must contain at least one node and **strictly fewer nodes -than the matched pattern**. Equal-size rewrites and expansions are rejected, -including dynamic recipes, retaining the node-count termination policy. -All internal references must point backward to valid output slots; disconnected -nodes and undeclared domains are errors. Limits remain 64 nodes and 64 inputs/ -outputs per node. Pattern parsing additionally bounds named values to 256. - -All replacement nodes are staged and their imported schemas/attributes checked -before deleting matched nodes or rewiring edges. A schema failure removes every -staged node and synthesized initializer, leaving the original live graph intact; -unused staging `NodeArg`s and consumed unique-name IDs may remain. Final Resolve -is still responsible for type/shape inference and validation; arbitrary computed -recipes must establish semantic/type compatibility, and post-commit Resolve -failure is not a general graph transaction rollback. -All added nodes inherit root provider/layer metadata. Producer/consumer maps, -internal edges and external root edges are updated; every added node seeds the -invalidation worklist. Replacement notification attributes the removed pattern -to the final node. - -### Multiple outputs - -`MultiOp(name, inputs, output_count, domain)` returns an ordered vector of distinct -value handles. `domain` is optional, and an `Operator` overload is also available. -For example: +The following rule folds division by a finite float32 scalar into FusedMatMul's +`alpha` attribute. It requires the Microsoft-domain opset to be imported already. ```cpp -Pattern p; -auto x = p.Input(); -auto pieces = p.MultiOp("Split", {x}, 2); -p.Op("Concat", {pieces[0], pieces[1]}); +#include +#include +#include "core/graph/graph_utils.h" +#include "core/optimizer/initializer.h" + +std::optional FloatScalar(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(); + return std::isfinite(scalar) ? std::optional(scalar) : std::nullopt; +} + +Rule MakeMatMulDivRule() { + Pattern pattern("(x,y,divisor) => mm = MatMul(x,y); return Div(mm,divisor)"); + auto guard = [](const Graph& graph, const Match& match) { + auto divisor = FloatScalar(graph, *match.valueOf("divisor")); + return IsFloat32(*match.valueOf("x")) && IsFloat32(*match.valueOf("y")) && + divisor && *divisor != 0.0f && std::isfinite(1.0f / *divisor); + }; + Replacement plan("(x,y,alpha) => com.microsoft::FusedMatMul(x,y,alpha=alpha)"); + auto replacement = [plan](const Graph& graph, const Match& match) { + ONNX_NAMESPACE::AttributeProto alpha; + alpha.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); + alpha.set_f(1.0f / *FloatScalar(graph, *match.valueOf("divisor"))); + match.update("alpha", std::move(alpha)); + return plan; + }; + return Rule{"MatMulDiv", std::move(pattern), + {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}, guard, replacement}; +} ``` -The output count must match the node exactly. Matching can find a producer -through any output, including its second or later output, and checks every -bound output against its exact output slot. Unused outputs of a connected node -are allowed in the definition, but cannot be graph outputs or have external -consumers if the node is deleted. The tests include Split/Concat and a genuine -ONNXScript LayerNorm-bias adaptation with unused mean/inverse-standard-deviation -outputs; exposing either statistic prevents its deletion. - -A multi-output root is also supported. The final replacement node inherits -**all root outputs in the original order**, including their names, types, shapes, -graph-output status, and consumers. The recipe does not need an output list. -Its schema must permit that output count. Identity/Split to Split tests execute -both root outputs and downstream fanout, including repeated uses of output 1. -This does not yet support mapping observable outputs from multiple matched nodes -to arbitrary replacement slots, or matching omitted optional output placeholders. - -### Guards, applicability and integration - -Guards and recipes are trusted C++ code: they must be pure, guards must depend on -local bindings/attributes/constants rather than mutable global graph conditions, -they must not retain match -pointers, and must establish semantic equivalence and replacement type -compatibility. The caller must establish kernel compatibility for its EP policy. -The engine provides structural and deletion safety; it cannot prove -an arbitrary callback's mathematical correctness. Guard false means an ordinary -nonmatch; invalid patterns or recipes report explicit errors. - -Each `OpsetRequirement` is `{domain, min_version, max_version}`, with inclusive, -positive bounds. A rule must declare every domain used by its pattern or -replacement; additional domain requirements are allowed. All requirements must -be satisfied by the graph's import map. Missing imports or out-of-range versions -make the whole rule inapplicable, before matching or guards. Duplicate domains -(including `""`/`"ai.onnx"` aliases), reversed ranges and undeclared pattern -domains are rejected when compiling rules. - -For example, MatMul/Div to FusedMatMul declares -`{{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}`. The standard-domain nodes use the -model's ONNX import; the replacement uses the existing Microsoft-domain import. -Rule authors are responsible for equivalence throughout the declared ranges. -The engine neither derives nor compares node schema since-versions. Replacement -schemas are resolved normally from imports and validated before deletion. - -`Apply` recurses into subgraphs and resolves the graph through the normal ORT -transformer wrapper. For per-call statistics and explicit graph resolution: +This example uses `IsFloat32` from the first example. The callback is called only +after its guard succeeds. Guards must be pure; neither guard nor callback may +mutate the graph, retain match pointers, or rely on mutable global graph state. +Use local bindings, attributes, and constants to establish applicability. -```cpp -PatternRewriter rewriter(MakeExampleRules()); -Statistics stats; -ORT_RETURN_IF_ERROR(rewriter.Run(graph, stats)); // already-resolved graph, one graph only -ORT_RETURN_IF_ERROR(graph.Resolve()); // required after rewriting -``` +`update` accepts a concrete `AttributeProto`, including list-valued attributes. +It cannot overwrite a named pattern value or store a referenced/undefined +attribute. Extra bindings are cleared before each replacement callback; +`attributeOf` retrieves one 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. + +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. + +## When a match succeeds or fails + +Before calling a rule's guard, the matcher checks: + +- Exact operator name/domain, ordered input/output counts, and repeated binding + identity. Missing optional slots and 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 non-root output stays internal. A graph output, an external consumer, or + an implicit capture by an external subgraph prevents the rewrite. +- Execution-provider eligibility as described below. -Callers must not concurrently mutate a graph. Compiled rules and options are -immutable and per-call state/statistics are local. +The root may have graph outputs or external consumers, because its complete +output list is preserved. Even match-only mode applies deletion-boundary checks: +it reports safe rewrite eligibility, not arbitrary structural occurrences. -### Generic execution-provider policy +### Types in the supplied examples -Rules do not inspect EP assignments. After structural matching and before rule -guards, the engine applies `CanRewritePartition`: all matched nodes must have -the same assignment, and an assigned EP must belong to -`Options::compatible_execution_providers`. Unassigned nodes are always eligible. -The default allowed set contains only CPU, preserving the verified execution scope. -Callers can opt into an EP (or clear the set for unassigned-only rewrites): +`MakeExampleRules()` supplies six rules when contrib operators are enabled: + +| 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 guard 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 builds supply only the two HardSwish rules. + +## 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 guard 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 guard and +also in match-only mode. To opt into another provider: ```cpp Options options; options.compatible_execution_providers.insert(kCudaExecutionProvider); -PatternRewriter rewriter(MakeExampleRules(), options); +PatternRewriter rewriter(Rules{MakeHardSwishRule()}, options); ``` -The replacement inherits the root's assignment. Allowing an EP is a caller assertion -of kernel/precision compatibility, not automatic kernel discovery or evidence of -GPU execution. Mixed assignments are rejected even if both EPs are allowed. This -generic policy also applies in match-only mode, which measures safe rewrite -eligibility rather than partition-independent structural matches. - -## Compilation, matching, mutation, and termination - -Construction validates connected acyclic definitions, used inputs, rule-level -domain/opset requirements, and limits of 64 nodes, 256 values and 64 inputs/outputs per operator. -It buckets groups by root op/domain and interns identical structural signatures, -including ordered output handles, binding identities and scalar literal type/value -constraints, but not diagnostic names or rule-level import ranges. -Rules with identical topology and literal constraints can share one structural -match even when their guards or applicability ranges differ. Patterns containing -attribute predicates currently receive separate plans, including equality predicates: -opaque `std::function` callbacks cannot safely be compared. This conservative -choice preserves correctness at the cost of repeated matching for attribute-heavy -rule sets; arbitrary callbacks also have no engine-imposed time bound. -There is **no arbitrary-prefix trie**: different structural -plans still match independently. This is intentionally a simpler, measurable -starting point than a decision DAG compiler. - -Each call filters rules once against `Graph::DomainToVersionMap()`, retaining -registration order within each structural group. No applicability state is cached -across graphs; recursive subgraph calls use their own graph import context. -Filtering costs the sum of rule requirement counts, outside the node traversal. -If no rules apply, the call returns without traversing nodes. Model-local function -rewriting remains unsupported; future function rewriting must use that function's -imports, not the enclosing model's. - -The call then builds a producer index and topological worklist once. Matching walks -the compiled node definitions backward without recursive search, alternation, -commutative backtracking, or unbounded neighborhood exploration. Each node checks -op/domain, exact input/output counts, repeated bindings and layer annotation -agreement, plus any attribute predicates and scalar literal constraints. -Provider eligibility is a separate generic check. All non-root outputs -must be internal: graph outputs, -external consumers (including implicit subgraph captures), and pattern inputs -produced inside the deletion set prevent replacement. - -Registration order defines priority even across different structural groups. -Match-only mode reports all passing rule guards; rewrite mode selects only the -first passing rule. Replacements may consume **only matched input variables, -new scalar literals or earlier replacement outputs**, retain the -root's output `NodeArg`s (name/type/shape and external users), and inherit its -execution provider and layer annotation. The replacement domain must be declared -in the rule requirements and already imported for every added node. Replacement schema/attribute checks -precede deletion. ORT edges and producer/consumer indexes are updated; removed -nodes are reported via `NotifyNodeReplacement`. Final Resolve validates types and -inferred shapes, rather than silently ignoring graph-validation errors. - -After mutation, all replacement nodes and surviving input producers seed a bounded -forward invalidation walk. Descendant roots up to the maximum pattern size are -requeued, since producer changes or reduced use counts can unlock deeper patterns. -Queue deduplication avoids redundant pending entries; removed IDs are skipped. -Traversal and rule priority are deterministic. - -Every rule replaces at least two nodes with a smaller nonempty node sequence. Thus every -successful rewrite strictly decreases live node count, and a graph with N nodes -admits at most N-1 rewrites. This is a termination policy, **not a runtime bound**: -guards, fanout, structural checks, and worklist invalidation still cost time. -Expansion rules and one-node rewrites are deliberately rejected. - -For match-only workloads, let U be the number of distinct structural plans in a -root bucket and P the maximum pattern size. Structure work is bounded by the -sum over visited roots of U times pattern/boundary checking cost, plus passing -plan guard costs. Boundary checks include matched-node fanout; they are not -constant-time. The reference mode performs the same checks per rule instead, -with cheap root op/domain rejection before allocating binding buffers. -Rewrite workloads additionally pay edge updates and bounded descendant walks; -large fanout can make these expensive. Arbitrary trusted guards have no imposed -time bound. These structural bounds do not provide mathematical worst-case runtime -guarantees for arbitrary rules or models. - -`visited` counts live worklist roots, including revisits. `candidates` counts -enumerated rule IDs (all rules in the reference mode, root-indexed rules in the -indexed mode); `structural_checks` counts plan attempts, including cheap root -rejections in the reference. `guards` counts semantic callback evaluations after -structural and deletion-boundary checks. `matches` counts all passing guards in -match-only mode, or selected winners in rewrite mode; `rewrites` counts completed -mutations. Match-only therefore measures safe rewrite eligibility, not a search -that ignores external consumers. - -## Intentional limitations and precision - -Matched and replacement nodes may have multiple outputs, but the final replacement -node must preserve the root's complete ordered output list. Optional/missing -inputs or outputs, variadic patterns, -model-local function rewriting, general tensor synthesis, domain-import edits, alternative -patterns and implicit commutation are unsupported. Control-flow nodes are not -matched, but `Apply` rewrites inside their subgraphs. External tensor data and -overridable initializers do not satisfy example constant guards; scalar constants -must be finite float32 rank-zero initializers. Singleton vectors are deliberately -not treated as scalars, since broadcasting can otherwise change the output rank. -Constant nodes require an earlier constant-folding pass. - -The examples support a conservative subset, not every ONNXScript variant: -RMSNorm requires the last axis, keepdims, positive finite epsilon, float32, -known positive-width matching one-dimensional scale, and pre-opset-18 attribute -axes. Transpose requires known rank >=2 and exactly the last-two-axis swap. -RMSNorm declares ONNX opsets 7 through 17. MatMul Div requires a finite nonzero -scalar and finite reciprocal. The FusedMatMul examples declare ONNX opsets 9 -through 21 and require an existing `com.microsoft` opset-1 import; they do not add -one. The upper bounds are conservative supported ranges, not inferred schema -version boundaries; extending them requires reviewing and testing semantics. -HardSwish requires exact float32 constants/attributes and model opset 14 through -21. The tested CPU build executes HardSwish by expanding its ONNX function body -to HardSigmoid/Mul; RMSNorm and FusedMatMul execute native CPU kernels. -Floating-point fusions change evaluation order; equivalence means tested -tolerance, not bitwise identity. NaNs/infinities and every extreme finite input -are not claimed to have bitwise-equivalent behavior. - -## Correctness tests - -The normal tests live in `onnxruntime/test/optimizer/pattern_rewriter_test.cc`. -They include executed CPU equivalence and structural assertions, not only match -counts. Use an isolated build directory and virtual environment: - -```powershell -python -m venv .venv -.\.venv\Scripts\python.exe tools\ci_build\build.py --build_dir build\rewriter --config Release --update --parallel 4 --skip_tests --cmake_generator "Visual Studio 17 2022" --cmake_extra_defines FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER -.\.venv\Scripts\python.exe tools\ci_build\build.py --build_dir build\rewriter --config Release --build --parallel 4 --target onnxruntime_test_all --skip_tests -Set-Location build\rewriter\Release\Release -.\onnxruntime_test_all.exe --gtest_filter="PatternRewriterTests.*" +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. + +## Applying rules, priority, and errors + +Register the transformer through the normal session pipeline: + +```cpp +auto transformer = std::make_unique(Rules{MakeHardSwishRule()}); +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{MakeHardSwishRule()}); +Statistics statistics; +ORT_RETURN_IF_ERROR(rewriter.Run(graph, statistics)); +ORT_RETURN_IF_ERROR(graph.Resolve()); ``` -Ensure CMake is on PATH (for example, Visual Studio's bundled CMake). -`FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER` avoids accidentally importing an -incompatible installed dependency, such as a Debug-only protobuf installation -while building Release. Dependencies are built inside this isolated -build tree instead of modifying the shared installation. - -## Decoder correctness topology - -`BuildDecoder` is a **hand-constructed, model-derived eager Llama-style topology**, -not an exported pretrained model or a claim of identical exporter node order. -Sources at ORT base revision -[`42518d89d7c5393baf97d69f8c2670668c40decf`](https://github.com/microsoft/onnxruntime/tree/42518d89d7c5393baf97d69f8c2670668c40decf) -are `onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py` -(eager decoder, dynamic past/present KV exports), -`models/llama/llama_inputs.py` (cache layout), -`onnxruntime/test/optimizer/group_query_attention_pre_norm_fusion_test.cc` -(small Q/K/V/head/cache geometry), and ONNXScript's RMSNorm source above. - -Each block has two decomposed RMSNorms, Q/K/V projections, reshape/transposes, -split-half RoPE via Slice/Neg/Concat and gathered sine/cosine position caches, -past/current KV concatenation with per-layer graph outputs, grouped-head expansion, -QK transpose/MatMul/scaling/causal-mask/Softmax/AV attention, output projection, -two residual adds, and a SiLU-gated up/gate/down MLP. B=1, query heads=4, -KV heads=2, head size=4, hidden width=16, MLP width=64, past length=3. -Sine/cosine caches use standard paired frequencies; weights are tiny random -float32 tensors. No tokenizer, embedding lookup, LM head, dynamic position/cache -control flow, mixed precision, quantization or production tensor sizes are modeled. -Both sequence=1 decode and sequence=4 causal prefill have executed equivalence tests. +`Run` handles one graph. The transformer's `Apply` path recurses into subgraphs +and performs normal graph resolution. Callers must not concurrently mutate a +graph. Options and compiled rules are immutable; match state and statistics are +per call. + +Rules have deterministic **registration-order priority**, including across +different structural groups. At each eligible root, rewriting chooses the first +passing guard. With `Options{.match_only=true}`, all passing guards are counted +and replacements are not called. `indexed=false` selects the reference matching +mode with the same eligibility and priority semantics. + +Every pattern must contain at least two nodes. A replacement must have a +nonempty node sequence with **strictly fewer nodes** than the pattern, and its +final output count must equal the root's output count. Expansion and one-node +rewrites are rejected. Each successful rewrite reduces live node count, so a +graph with N nodes admits at most N-1 rewrites. This establishes termination, +not a time bound for arbitrary callbacks or graph fanout. + +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 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. Guard false +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, guard calls, successful matches, and completed +rewrites. Match-only counts all passing guards; rewrite mode counts selected +winners. These counters do not alter rule semantics. + +## Limits and unsupported forms + +Definitions are limited to 64 nodes, 256 values, and 64 inputs/outputs per +operator. They are connected acyclic graphs rooted at the final operator. +There is no alternation, implicit commutation, optional/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, +guards, 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. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 30ffc78bfc0c8..afafd9365fd52 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -130,15 +130,12 @@ bool CanRewritePartition(const Match& match, const Options& options) { Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const Replacement& recipe, Producers& producers, InlinedVector& added) { - ORT_RETURN_IF(recipe.Plan() && (!recipe.op.empty() || !recipe.domain.empty() || !recipe.inputs.empty() || - !recipe.attributes.empty()), - "A replacement plan cannot be combined with single-node recipe fields."); - InlinedVector single; - if (!recipe.Plan()) single.push_back({recipe.op, recipe.domain, recipe.inputs, recipe.attributes, 1}); - const auto& nodes = recipe.Plan() ? *recipe.Plan() : single; + const auto& nodes = *recipe.Plan(); ORT_RETURN_IF(nodes.empty() || nodes.size() >= match.nodes.size(), "Replacements must contain fewer nodes than the matched pattern."); const Node& root = *match.nodes.back(); + ORT_RETURN_IF(nodes.back().output_count != root.OutputDefs().size(), + "The final replacement must preserve the root output count."); InlinedVector> external_inputs; for (size_t i = 0; i < nodes.size(); ++i) { const auto& node = nodes[i]; @@ -169,14 +166,7 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R if (const auto* handle = std::get_if(&input)) { value = *handle; } else { - const auto* argument = std::holds_alternative(input) ? match.valueOf(std::get(input)) - : std::get(input); - for (int id = 0; id < static_cast(match.values.size()); ++id) { - if (argument && match.values[id] == argument && pattern.IsInput(id)) { - value = id; - break; - } - } + value = pattern.NamedValue(std::get(input)); } ORT_RETURN_IF( value.handle < 0 || static_cast(value.handle) >= match.values.size() || !pattern.IsInput(value), @@ -372,10 +362,6 @@ void Pattern::Name(Value value, std::string_view name) { } } -void Pattern::Attributes(std::initializer_list matchers) { - Attributes(InlinedVector(matchers)); -} - void Pattern::Attributes(InlinedVector matchers) { ORT_ENFORCE(!nodes_.empty(), "Attribute matchers require a pattern node."); InlinedHashSet names; @@ -388,12 +374,6 @@ void Pattern::Attributes(InlinedVector matchers) { target.insert(target.end(), std::make_move_iterator(matchers.begin()), std::make_move_iterator(matchers.end())); } -Value Pattern::Op(Operator op, std::initializer_list inputs) { return MultiOp(std::move(op), inputs, 1)[0]; } - -InlinedVector Pattern::MultiOp(Operator op, std::initializer_list inputs, size_t output_count) { - return MultiOp(std::move(op), gsl::span(inputs.begin(), inputs.size()), output_count); -} - 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."); @@ -419,7 +399,7 @@ bool Pattern::IsInput(Value value) const { }); } -namespace { +namespace detail { class PatternParser { public: explicit PatternParser(std::string_view text, InlinedVector* attributes = nullptr, @@ -452,17 +432,19 @@ class PatternParser { if (kind_ == Kind::Identifier && token_ == "return") { returned = true; ORT_RETURN_IF_ERROR(Next()); - } else { + } + { const size_t start = token_start_; do { std::string name; - ORT_RETURN_IF_ERROR(Identifier(name, replacement_attributes_ && names.empty())); + ORT_RETURN_IF_ERROR(Identifier(name, names.empty())); ORT_RETURN_IF(std::find(names.begin(), names.end(), name) != names.end(), Error("duplicate assignment name")); names.push_back(std::move(name)); if (token_ != ",") break; ORT_RETURN_IF_ERROR(Next()); } while (true); - if (replacement_attributes_ && names.size() == 1 && (token_ == "(" || token_ == "::")) { + 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(); @@ -700,20 +682,15 @@ class PatternParser { InlinedVector, 1>* replacement_bindings_ = nullptr; InlinedHashSet attribute_inputs_; }; -} // namespace - -Pattern::Pattern(std::string_view text) { ORT_THROW_IF_ERROR(PatternParser(text).Parse(*this)); } +} // namespace detail -Replacement::Replacement(InlinedVector nodes) { - ORT_ENFORCE(!nodes.empty() && nodes.size() <= 64, "Replacements need between 1 and 64 nodes."); - plan_ = std::make_shared>(std::move(nodes)); -} +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(PatternParser(text, &parsed_attributes, &parsed_bindings).Parse(parsed)); + ORT_THROW_IF_ERROR(detail::PatternParser(text, &parsed_attributes, &parsed_bindings).Parse(parsed)); InlinedHashMap outputs; InlinedVector nodes; InlinedHashSet attribute_names; @@ -934,17 +911,14 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { InlinedVector added; winning_match.ClearUpdates(); auto recipe = rules_[winner].replacement(graph, winning_match); + ORT_RETURN_IF(!recipe.Plan(), "Replacement has no parsed plan."); 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; }); }; - if (recipe.Plan()) { - for (const auto& node : *recipe.Plan()) - ORT_RETURN_IF(!declared(node.domain), "Replacement domain has no rule opset requirement: ", node.domain); - } else { - ORT_RETURN_IF(!declared(recipe.domain), "Replacement domain has no rule opset requirement: ", recipe.domain); - } + for (const auto& node : *recipe.Plan()) + ORT_RETURN_IF(!declared(node.domain), "Replacement domain has no rule opset requirement: ", node.domain); ORT_RETURN_IF_ERROR(Replace(graph, *winning_pattern, winning_match, recipe, producers, added)); ++statistics.rewrites; affected.insert(affected.end(), added.begin(), added.end()); diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 26284b7f09010..a01419fb7e31a 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -3,11 +3,9 @@ #pragma once #include -#include #include #include #include -#include #include #include @@ -16,6 +14,10 @@ namespace onnxruntime { namespace rewrite { +namespace detail { +class PatternParser; +} + struct Value { std::string name; int handle = -1; @@ -54,30 +56,8 @@ struct OpsetRequirement { // Reusing a Value requires exactly the same NodeArg, including for input variables. class Pattern { public: - Pattern() = default; explicit Pattern(std::string_view text); - Value Input(std::string_view name = ""); - Value Constant(float value); - Value Constant(int64_t value); - Value Constant(int value) { return Constant(int64_t{value}); } - Value Constant(double value) { return Constant(static_cast(value)); } Value NamedValue(std::string_view name) const; - void Name(Value value, std::string_view name); - void Attributes(std::initializer_list matchers); - void Attributes(InlinedVector matchers); - Value Op(Operator op, std::initializer_list inputs); - template , int> = 0> - Value Op(Name&& name, std::initializer_list inputs, std::string_view domain = "") { - return Op(Operator{std::string(std::string_view(name)), std::string(domain)}, inputs); - } - InlinedVector MultiOp(Operator op, std::initializer_list inputs, size_t output_count); - InlinedVector MultiOp(Operator op, gsl::span inputs, size_t output_count); - template , int> = 0> - InlinedVector MultiOp(Name&& name, std::initializer_list inputs, size_t output_count, - std::string_view domain = "") { - return MultiOp(Operator{std::string(std::string_view(name)), std::string(domain)}, inputs, output_count); - } - struct NodePattern { Operator op; InlinedVector inputs; @@ -90,8 +70,18 @@ class Pattern { const auto& Constants() const { return constants_; } private: + friend class detail::PatternParser; + 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); + 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; InlinedVector nodes_; InlinedVector values_; @@ -120,7 +110,7 @@ struct Replacement { struct Scalar { std::variant value; }; - using Input = std::variant; + using Input = std::variant; struct NodeRecipe { std::string op; std::string domain; @@ -129,21 +119,10 @@ struct Replacement { size_t output_count = 1; InlinedHashMap attribute_bindings; }; - Replacement() = default; - Replacement(std::string op_name, std::string op_domain, InlinedVector arguments, NodeAttributes attrs) - : op(std::move(op_name)), - domain(std::move(op_domain)), - inputs(std::move(arguments)), - attributes(std::move(attrs)) {} explicit Replacement(std::string_view text); - explicit Replacement(InlinedVector nodes); // 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_; } - std::string op; - std::string domain; - InlinedVector inputs; - NodeAttributes attributes; private: friend class PatternRewriter; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 6cd6b34d84aad..e09caad2e4ea2 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "core/optimizer/pattern_rewriter.h" @@ -76,21 +77,21 @@ size_t Count(const Graph& graph, const std::string& op) { Rule IdentityPair(const char* name = "identity", bool accept = true) { Pattern p("(x) => a = Identity(x); return Identity(a)"); auto guard = [accept](const Graph&, const Match&) { return accept; }; - auto replacement = [](const Graph&, const Match& m) { - return Replacement{"Identity", kOnnxDomain, {m.valueOf("x")}, {}}; - }; + Replacement replacement("(x) => Identity(x)"); Rule rule{name, p, {{kOnnxDomain, 1, 21}}, guard, replacement}; return rule; } Rule IdentitySplit() { - Pattern p; - auto x = p.Input(); - auto first = p.Op("Identity", {x}); - p.MultiOp("Split", {first}, 2); + Pattern p("(x) => first = Identity(x); return a, b = Split(first)"); auto guard = [](const Graph&, const Match&) { return true; }; - auto replacement = [x](const Graph&, const Match& match) { - return Replacement{"Split", kOnnxDomain, {x}, match.nodes.back()->GetAttributes()}; + Replacement plan("(x, axis) => return a, b = Split(x, axis=axis)"); + auto replacement = [plan](const Graph&, const 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 plan; }; Rule rule{"IdentitySplit", p, {{kOnnxDomain, 13, 17}}, guard, replacement}; return rule; @@ -104,9 +105,7 @@ Rule SplitConcat() { return (axis == split_attributes.end() || axis->second.i() == 0) && match.nodes[1]->GetAttributes().at("axis").i() == 0; }; - auto replacement = [](const Graph&, const Match& m) { - return Replacement{"Identity", kOnnxDomain, {m.valueOf("x")}, {}}; - }; + Replacement replacement("(x) => Identity(x)"); Rule rule{"SplitConcat", p, {{kOnnxDomain, 13, 17}}, guard, replacement}; return rule; } @@ -130,18 +129,24 @@ Rule LayerNormBias() { width.dim_value() == bias_width.dim_value() && (axis == match.nodes[0]->GetAttributes().end() || axis->second.i() == -1); }; - auto replacement = [](const Graph&, const Match& match) { - return Replacement{"LayerNormalization", - "", - {match.valueOf("x"), match.valueOf("scale"), match.valueOf("bias")}, - match.nodes[0]->GetAttributes()}; + Replacement plan(R"( + (x, scale, bias, axis, epsilon, stash_type) => + LayerNormalization(x, scale, bias, axis=axis, epsilon=epsilon, stash_type=stash_type) + )"); + auto replacement = [plan](const Graph&, const Match& match) { + 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 plan; }; Rule rule{"LayerNormBias", p, {{kOnnxDomain, 17, 17}}, guard, replacement}; return rule; } // Small-weight eager decoder: B=1, Hq=4, Hkv=2, D=4. Cache length is 3. -// See the design document for the model/export sources and deliberate simplifications. 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); @@ -290,25 +295,19 @@ TEST(PatternRewriterTests, HardSigmoidAttributeGuard) { } } -TEST(PatternRewriterTests, NamedValuesAndAttributePredicates) { - Pattern p; - auto x = p.Input("x"); - auto first = p.Op("HardSigmoid", {x}); - p.Name(first, "sigmoid"); - p.Attributes({AttributeMatcher::Equal("alpha", 0.25f), - {"beta", [](const auto& attr) { - return attr.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT && attr.f() > 0.4f; - }}}); - p.Op("Mul", {first, x}); +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(p.Input("x"), OnnxRuntimeException); + EXPECT_THROW(Pattern("(x,x) => return Identity(x)"), OnnxRuntimeException); EXPECT_THROW(p.NamedValue("missing"), OnnxRuntimeException); - EXPECT_THROW(p.Name(x, "sigmoid"), OnnxRuntimeException); - EXPECT_THROW(p.Op("Identity", {Value{-1}}), 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; @@ -333,6 +332,10 @@ TEST(PatternRewriterTests, NamedValuesAndAttributePredicates) { ASSERT_STATUS_OK(test.graph.Resolve()); auto rule = IdentityPair(); rule.pattern = p; + rule.guard = [](const Graph&, const 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); @@ -357,7 +360,7 @@ TEST(PatternRewriterTests, NamedMatchLookupAndReplacementSafety) { test.builder.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); auto first = IdentityPair(), second = IdentityPair(); - second.pattern = Pattern("(input) => middle = Identity(input); return Identity(middle)"); + second.pattern = Pattern("(input) => middle = Identity(input); return root = Identity(middle)"); first.guard = [=](const Graph&, const Match& m) { EXPECT_EQ(m.valueOf("x"), x); EXPECT_EQ(m.valueOf("a"), intermediate); @@ -371,12 +374,11 @@ TEST(PatternRewriterTests, NamedMatchLookupAndReplacementSafety) { EXPECT_THROW(m.valueOf("x"), OnnxRuntimeException); return true; }; - second.replacement = [scenario, output](const Graph&, const Match& m) { - const NodeArg* input = scenario == 0 ? m.valueOf("input") - : scenario == 1 ? m.valueOf("middle") - : scenario == 2 ? output - : nullptr; - return Replacement{"Identity", "", {input}, {}}; + 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; @@ -482,13 +484,8 @@ TEST(PatternRewriterTests, ReplacementBranchesAndMultiOutputRoot) { std::make_unique(Rules{branched})); auto split = IdentitySplit(); - Pattern p; - auto x = p.Input("x"); - auto a = p.Op("Identity", {x}); - auto b = p.Op("Identity", {a}); - p.MultiOp("Split", {b}, 2); - split.pattern = p; - split.replacement = Replacement("(x) => a=Identity(x); return Split(a, axis=0)"); + split.pattern = Pattern("(x) => a = Identity(x); b = Identity(a); return left, right = Split(b)"); + split.replacement = Replacement("(x) => a=Identity(x); return left, right = Split(a, axis=0)"); auto build_split = [](ModelTestBuilder& b) { auto* x = b.MakeInput({4}, -1.0f, 1.0f); auto* a = b.MakeIntermediate(); @@ -509,7 +506,10 @@ TEST(PatternRewriterTests, ReplacementBranchesAndMultiOutputRoot) { 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)=>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)"); @@ -522,42 +522,19 @@ TEST(PatternRewriterTests, ReplacementParsingCompilationAndRollback) { EXPECT_THROW(PatternRewriter({invalid}), OnnxRuntimeException); EXPECT_NO_THROW(Replacement("(Identity)=>Identity(Identity)")); for (bool indexed : {false, true}) { - for (int scenario = 0; scenario < 9; ++scenario) { + 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]; - InlinedVector nodes{{"Identity", "", {std::string("x")}, {}, 1}, - {"HardSwish", "", {Replacement::Output{0}}, {}, 1}}; - if (scenario == 0) nodes[1].op = "NoSuchReplacementOperator"; - if (scenario == 1) nodes[1].inputs = {Replacement::Output{1}}; - if (scenario == 2) nodes[1].inputs = {Replacement::Output{0, 1}}; - if (scenario == 3) nodes[1].inputs = {std::string("x")}; - if (scenario == 4) { - ONNX_NAMESPACE::AttributeProto attr; - attr.set_name("axis"); - attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); - attr.set_s("wrong"); - nodes[1].op = "Softmax"; - nodes[1].attributes.emplace("axis", attr); - } - if (scenario == 5) nodes[1].domain = kMSDomain; - if (scenario == 6) nodes[0].output_count = 0; - if (scenario == 7) { - nodes[0].op = "Add"; - nodes[0].inputs.push_back(Replacement::Scalar{0.0f}); - nodes[1].op = "NoSuchReplacementOperator"; - } - rule.replacement = [recipe = Replacement(std::move(nodes))](const Graph&, const Match&) { return recipe; }; - if (scenario == 8) { - rule.replacement = [](const Graph&, const Match&) { - Replacement recipe("(x)=>HardSwish(x)"); - recipe.op = "Identity"; - return recipe; - }; - } + 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); @@ -572,7 +549,7 @@ TEST(PatternRewriterTests, ReplacementParsingCompilationAndRollback) { } } -TEST(PatternRewriterTests, ReplacementScalarConstantsAndProgrammaticSequence) { +TEST(PatternRewriterTests, ReplacementScalarConstantsAndParsedSequence) { auto rule = IdentityPair(); rule.pattern = Pattern("(x)=>a=Identity(x);b=Identity(a);return Identity(b)"); for (bool integer : {false, true}) { @@ -599,8 +576,7 @@ TEST(PatternRewriterTests, ReplacementScalarConstantsAndProgrammaticSequence) { }, TransformerLevel::Default, TransformerLevel::Level1, 17, 0, 0, std::make_unique(Rules{rule})); } - rule.replacement = Replacement(InlinedVector{ - {"Identity", "", {std::string("x")}, {}, 1}, {"Identity", "", {Replacement::Output{0}}, {}, 1}}); + 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; @@ -1070,9 +1046,9 @@ TEST(PatternRewriterTests, LiteralTypeAndRankMatching) { EXPECT_EQ(stats.matches, 0u); } } - Pattern p; - auto integer = p.Constant(3); - auto floating = p.Constant(3.0); + 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))); } @@ -1210,12 +1186,9 @@ TEST(PatternRewriterTests, PriorityOverlapTerminationAndMatchOnly) { } TEST(PatternRewriterTests, InvalidRuleAndRecipe) { - Pattern disconnected; - auto x = disconnected.Input(); - disconnected.Op({"Identity", ""}, {x}); - disconnected.Op({"Identity", ""}, {x}); + EXPECT_THROW(Pattern("(x)=>a=Identity(x);return Identity(x)"), OnnxRuntimeException); auto rule = IdentityPair(); - rule.pattern = disconnected; + rule.pattern = Pattern("(x)=>return Identity(x)"); EXPECT_THROW(PatternRewriter({rule}), OnnxRuntimeException); TestGraph test; auto* input = test.builder.MakeInput({1}, -1.0f, 1.0f); @@ -1225,22 +1198,9 @@ TEST(PatternRewriterTests, InvalidRuleAndRecipe) { test.builder.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); rule = IdentityPair(); - for (int scenario = 0; scenario < 5; ++scenario) { - rule.replacement = [scenario](const Graph&, const Match&) { - Replacement recipe{"Identity", "", {0}, {}}; - if (scenario == 0) recipe.op = "NoSuchReplacementOperator"; - if (scenario == 1) recipe.domain = "not.imported"; - if (scenario == 2) recipe.inputs = {1}; // Deleted intermediate, not an external input. - if (scenario == 3) { - ONNX_NAMESPACE::AttributeProto attr; - attr.set_name("invalid"); - attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); - attr.set_i(1); - recipe.attributes.emplace("invalid", attr); - } - if (scenario == 4) recipe.domain = kMSDomain; // Imported, but not declared by this rule. - return recipe; - }; + 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); @@ -1266,17 +1226,13 @@ TEST(PatternRewriterTests, RuleImportRangesAndAliases) { for (bool indexed : {false, true}) { auto rule = IdentityPair(); rule.opsets = {{kOnnxDomainAlias, 14, 17}}; - Pattern pattern; - auto input = pattern.Input(); - auto first = pattern.Op({"Identity", kOnnxDomainAlias}, {input}); - pattern.Op({"Identity", kOnnxDomainAlias}, {first}); - rule.pattern = std::move(pattern); + rule.pattern = Pattern("(x)=>a=ai.onnx::Identity(x);return ai.onnx::Identity(a)"); int guards = 0; rule.guard = [&guards](const Graph&, const Match&) { ++guards; return true; }; - rule.replacement = [](const Graph&, const Match&) { return Replacement{"Identity", kOnnxDomainAlias, {0}, {}}; }; + 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}) { @@ -1404,13 +1360,9 @@ TEST(PatternRewriterTests, PriorityAcrossInterleavedStructuralGroups) { lower_priority.guard = [](const Graph& graph, const Match& m) { return graph.NodeProducesGraphOutput(*m.nodes.back()); }; - Pattern p; - auto input = p.Input(); - auto first = p.Op({"Identity", ""}, {input}); - auto second = p.Op({"Identity", ""}, {first}); - p.Op({"Identity", ""}, {second}); + Pattern p("(x)=>a=Identity(x);b=Identity(a);return Identity(b)"); auto guard = [](const Graph&, const Match&) { return true; }; - auto replacement = [input](const Graph&, const Match&) { return Replacement{"Identity", "", {input}, {}}; }; + Replacement replacement("(x)=>Identity(x)"); Rule higher_priority{"higher_priority", p, {{kOnnxDomain, 1, 21}}, guard, replacement}; Statistics stats; ASSERT_STATUS_OK( @@ -1457,10 +1409,7 @@ TEST(PatternRewriterTests, OutputCountMismatch) { b.AddNode("Dropout", {a}, {b.MakeOutput(), b.MakeOutput()}); b.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); - Pattern p; - auto input = p.Input(); - auto identity = p.Op({"Identity", ""}, {input}); - p.Op({"Dropout", ""}, {identity}); + Pattern p("(x)=>a=Identity(x);return Dropout(a)"); auto rule = IdentityPair(); rule.pattern = p; Statistics stats; @@ -1468,13 +1417,19 @@ TEST(PatternRewriterTests, OutputCountMismatch) { EXPECT_EQ(stats.rewrites, 0u); } -TEST(PatternRewriterTests, MultiOpDefinitionsAndOverloads) { - Pattern p; - auto x = p.Input(); - std::string identity = "Identity"; - auto a = p.Op(identity, {x}); - auto b = p.Op(std::string_view(identity), {a}, kOnnxDomainAlias); - auto outputs = p.MultiOp({"Split", kOnnxDomainAlias}, {b}, 2); +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);return result=Identity(a)")); + EXPECT_NO_THROW(Replacement("(x)=>return result=Identity(x)")); +} + +TEST(PatternRewriterTests, ParsedMultiOutputDefinitionsAndLimits) { + Pattern p("(x)=>a=Identity(x);b=ai.onnx::Identity(a);return left,right = ai.onnx::Split(b)"); + 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)); @@ -1482,12 +1437,28 @@ TEST(PatternRewriterTests, MultiOpDefinitionsAndOverloads) { EXPECT_FALSE(p.IsInput(outputs[1])); EXPECT_EQ(p.Nodes().back().op.domain, kOnnxDomain); const auto count = p.ValueCount(); - EXPECT_THROW(p.MultiOp("Split", {b}, 0), OnnxRuntimeException); - EXPECT_THROW(p.MultiOp("Split", {b}, 65), OnnxRuntimeException); - EXPECT_THROW(p.MultiOp("Split", {999}, 2), OnnxRuntimeException); + EXPECT_THROW(p = Pattern("(x)=>return = Split(x)"), OnnxRuntimeException); + EXPECT_THROW(p = Pattern("(x)=>return a,b = Split(missing)"), OnnxRuntimeException); + EXPECT_THROW(p = Pattern("(x)=>return left,left = Split(x)"), OnnxRuntimeException); + EXPECT_THROW(p = Pattern("(x)=>return x,right = Split(x)"), OnnxRuntimeException); + EXPECT_THROW(p = Pattern("(x)=>a,b=Split(x);return a,b"), OnnxRuntimeException); EXPECT_EQ(p.ValueCount(), count); - for (int i = 0; i < 3; ++i) p.MultiOp("Split", {b}, 64); - EXPECT_THROW(p.MultiOp("Split", {b}, 64), OnnxRuntimeException); + 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)=>return " + outputs_text(0, 65) + "=Split(x)"), 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 + "return " + outputs_text(3, 64) + "=Split(v2_0)"), OnnxRuntimeException); } TEST(PatternRewriterTests, MultiOutputIntermediateBoundariesAndBindings) { @@ -1529,10 +1500,7 @@ TEST(PatternRewriterTests, MultiOutputProducerReachedThroughSecondOutput) { test.builder.AddNode("Identity", {b}, {test.builder.MakeOutput()}); test.builder.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); - Pattern pattern; - auto input = pattern.Input(); - auto outputs = pattern.MultiOp("Split", {input}, 2); - pattern.Op("Identity", {outputs[1]}); + Pattern pattern("(x)=>a,b=Split(x);return Identity(b)"); auto rule = SplitConcat(); rule.pattern = pattern; rule.guard = [](const Graph&, const Match&) { return true; }; @@ -1557,19 +1525,23 @@ TEST(PatternRewriterTests, MultiOutputRootEdgesAndInvalidRecipe) { test.builder.AddNode("Add", {b, b}, {test.builder.MakeOutput()}); test.builder.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); - auto invalid = IdentitySplit(); - invalid.replacement = [](const Graph&, const Match&) { return Replacement{"Identity", "", {0}, {}}; }; Statistics stats; - EXPECT_FALSE(PatternRewriter({invalid}, {.indexed = indexed}).Run(test.graph, stats).IsOK()); - EXPECT_EQ(test.graph.NumberOfNodes(), 4); - ASSERT_STATUS_OK(test.graph.Resolve()); + for (auto text : {"(x)=>Identity(x)", "(x)=>Split(x)", "(x)=>return left,middle,right=Split(x)"}) { + 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"); - Pattern single_output; - auto input = single_output.Input(); - auto first = single_output.Op("Identity", {input}); - single_output.Op("Split", {first}); - wrong_count.pattern = single_output; - ASSERT_STATUS_OK(PatternRewriter({wrong_count, IdentitySplit()}, {.indexed = indexed}).Run(test.graph, stats)); + wrong_count.pattern = Pattern("(x)=>a=Identity(x);return Split(a)"); + auto valid = IdentitySplit(); + valid.guard = [a, b](const Graph&, const Match& match) { + EXPECT_EQ(match.valueOf("a"), a); + EXPECT_EQ(match.valueOf("b"), b); + return true; + }; + 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); @@ -1773,7 +1745,7 @@ TEST(PatternRewriterTests, ExampleGuardsCheckRelevantBindings) { } 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}, + : std::vector{4, 3}, -1.0f, 1.0f); bindings.push_back("y"); if (index == 3) { @@ -1829,13 +1801,9 @@ TEST(PatternRewriterTests, BranchedSharedNodeBindings) { if (external) b.AddNode("Identity", {sum}, {b.MakeOutput()}); b.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); - Pattern p; - auto px = p.Input(), py = p.Input(); - auto psum = p.Op({"Add", ""}, {px, py}); - auto pproduct = p.Op({"Mul", ""}, {psum, px}); - p.Op({"Add", ""}, {psum, pproduct}); + Pattern p("(x,y)=>sum=Add(x,y);product=Mul(sum,x);return Add(sum,product)"); auto guard = [](const Graph&, const Match&) { return true; }; - auto replacement = [](const Graph&, const Match&) { return Replacement{"Identity", "", {0}, {}}; }; + Replacement replacement("(x)=>Identity(x)"); Rule rule{"diamond", p, {{kOnnxDomain, 7, 21}}, guard, replacement}; Statistics stats; // Structural-only test: this deliberately has no semantically equivalent replacement. From f0ff67fe871aaaf7d52c6ed0f762115db6849fef Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 13 Oct 2026 18:17:25 -0700 Subject: [PATCH 04/45] Express extra-attribute rollback coverage with a parsed replacement Replace the last mutable single-node recipe use with a parsed plan containing an invalid schema attribute, preserving runtime failure and unchanged-graph assertions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/optimizer/pattern_rewriter_test.cc | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index e09caad2e4ea2..1d1aa528c80ed 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -731,9 +731,11 @@ TEST(PatternRewriterTests, ReplacementExtraAttributeBindings) { alpha_value.set_f(1.0f / 6.0f); m.update("alpha", alpha_value); if (scenario == 4) { - auto malformed = plan; - malformed.attributes.emplace("extra", beta); - return malformed; + return Replacement(R"( + (x,alpha,beta) => + h=HardSigmoid(x,alpha=alpha,beta=beta,extra=0.0) + return Mul(h,x) + )"); } return plan; }; From 8f15fdfab0b5af470f63d3506cd3bb3449f26a76 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 13 Oct 2026 18:19:46 -0700 Subject: [PATCH 05/45] Resolve callback replacement names against the winning rule Keep the winning rule's pattern rather than the shared structural group's pattern when lowering parsed callback replacements. This preserves rule-local names for shared plans and the existing named-lookup safety regression. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/optimizer/pattern_rewriter.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index afafd9365fd52..564b3e35726be 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -879,7 +879,7 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { if (id < winner) { winner = id; winning_match = match; - winning_pattern = &pattern; + winning_pattern = match.pattern; } if (!options_.match_only) break; } From 7d3850bac5e1386da4ace9f85a90ddbe7affc374 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 13 Oct 2026 18:22:32 -0700 Subject: [PATCH 06/45] Own the untyped guard-test value outside graph cleanup Graph::Resolve deletes unused untyped graph NodeArgs. Keep the negative guard-test NodeArg on the stack so copied matches never reference freed graph-owned storage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/optimizer/pattern_rewriter_test.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 1d1aa528c80ed..e27ef512f0905 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1726,7 +1726,8 @@ TEST(PatternRewriterTests, ExampleGuardsCheckRelevantBindings) { auto& b = test.builder; auto* non_float = b.MakeScalarInitializer(2.0); auto* integer = b.MakeScalarInitializer(2); - auto* unknown = b.MakeIntermediate(); + // 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"}; @@ -1772,7 +1773,7 @@ TEST(PatternRewriterTests, ExampleGuardsCheckRelevantBindings) { for (auto binding : bindings) { SCOPED_TRACE(binding); // Exercise the guard independently without making the resolved graph ill-typed. - for (const NodeArg* value : {non_float, integer, unknown}) { + for (const NodeArg* value : {non_float, integer, &unknown}) { Match changed = match; changed.values[match.pattern->NamedValue(binding).handle] = value; EXPECT_FALSE(guard(graph, changed)); From 1f386911c39692fbd1fb1ce02cb3c6f2fc45083e Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 13 Oct 2026 22:00:59 -0700 Subject: [PATCH 07/45] Separate named returns and support braced rewrite definitions Accept optional braced bodies and returns of all final operator outputs in order. Reject combined assignment-return syntax and migrate multi-output definitions and examples while preserving direct calls and unbraced compatibility. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 53 +++++++++------ .../core/optimizer/pattern_rewriter.cc | 28 ++++++-- .../core/optimizer/pattern_rewriter_rules.cc | 15 +++-- .../test/optimizer/pattern_rewriter_test.cc | 67 ++++++++++++++----- 4 files changed, 118 insertions(+), 45 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 2b6568f242fd3..74d21f5174145 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -30,9 +30,10 @@ bool IsFloat32(const NodeArg& value) { Rule MakeHardSwishRule() { Pattern pattern(R"( - (x) => + (x) => { sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) return Mul(sigmoid, x) + } )"); auto guard = [](const Graph&, const Match& match) { return IsFloat32(*match.valueOf("x")); @@ -54,31 +55,36 @@ change evaluation order; validate equivalence within the intended tolerance. ## Pattern and replacement syntax -Definitions start with named inputs and `=>`. Assignments create named operator -outputs; `return` identifies the final operator, called the root. Whitespace -separates calls, and semicolons are optional separators. Comments are not supported -inside definitions. +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 all +outputs of the final assignment. Whitespace separates calls, and semicolons are +optional separators. A trailing semicolon before the closing brace is allowed. +Nothing may follow the closing brace. Comments are not supported inside definitions. ```cpp Pattern pattern(R"( - (x) => + (x) => { add = Add(x, 3.0) clip = Clip(add, 0.0, 6.0) product = Mul(clip, x) - return result = Div(product, 6.0) + result = Div(product, 6.0) + return result; + } )"); Replacement replacement("(x) => HardSwish(x)"); ``` -The example has one root output, named `result`. A final call without an -assignment, such as `return Div(product, 6.0)`, also describes one output, +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(...)` for qualified operators: ```cpp -Pattern pattern("(x,y) => mm = com.microsoft::FusedMatMul(x,y); return Identity(mm)"); +Pattern pattern("(x,y) => { mm = com.microsoft::FusedMatMul(x,y); return Identity(mm); }"); Replacement replacement("(x,y) => com.microsoft::FusedMatMul(x,y)"); ``` @@ -106,7 +112,7 @@ an earlier constant-folding pass. Keyword arguments constrain attributes, not tensor inputs: ```cpp -Pattern pattern("(x) => h = HardSigmoid(x, alpha=0.25, beta=0.5); return Mul(h,x)"); +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. @@ -129,10 +135,11 @@ Assignments can describe a DAG rather than only a chain: ```cpp Pattern pattern(R"( - (x,y) => + (x,y) => { sum = Add(x,y) product = Mul(sum,x) return Add(sum,product) + } )"); ``` @@ -142,23 +149,31 @@ 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)"); +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, put the names on the final return assignment: +For a multi-output root, assign its outputs first, then return their names: ```cpp -Pattern pattern("(x) => copied = Identity(x); return left,right = Split(copied,axis=0)"); -Replacement replacement("(x) => return left,right = Split(x,axis=0)"); +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 root's two original output `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 only one output and cannot replace a two-output root. -Returning an existing value or a subset/reordering of an earlier operator's -outputs is not supported; the return must define the final operator. +A named return must list **all outputs of the final defined operator in their +original order**. Returning a leaf input, an earlier operator's outputs, a subset, +duplicate names, or reordered outputs is rejected. A named return does not create +an additional operator or introduce output selection. Intermediate operators may have unused outputs, but **none** of their outputs may be exposed outside the matched deletion set. @@ -191,7 +206,7 @@ std::optional FloatScalar(const Graph& graph, const NodeArg& value) { } Rule MakeMatMulDivRule() { - Pattern pattern("(x,y,divisor) => mm = MatMul(x,y); return Div(mm,divisor)"); + Pattern pattern("(x,y,divisor) => { mm = MatMul(x,y); return Div(mm,divisor); }"); auto guard = [](const Graph& graph, const Match& match) { auto divisor = FloatScalar(graph, *match.valueOf("divisor")); return IsFloat32(*match.valueOf("x")) && IsFloat32(*match.valueOf("y")) && diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 564b3e35726be..eae1ef96f710c 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -422,8 +422,10 @@ class PatternParser { } 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) { + while (kind_ != Kind::End && token_ != "}") { if (kind_ == Kind::Symbol && token_ == ";") { ORT_RETURN_IF_ERROR(Next()); continue; @@ -433,6 +435,7 @@ class PatternParser { returned = true; ORT_RETURN_IF_ERROR(Next()); } + bool returned_values = false; { const size_t start = token_start_; do { @@ -450,16 +453,31 @@ class PatternParser { names.clear(); returned = true; } else { - for (const auto& name : names) ORT_RETURN_IF(symbols_.count(name), Error("duplicate assignment name")); ORT_RETURN_IF(std::any_of(names.begin(), names.end(), [](const auto& name) { return name.find('.') != std::string::npos; }), Error("expected value identifier")); - ORT_RETURN_IF_ERROR(Expect("=")); + if (returned) { + ORT_RETURN_IF(token_ == "=", Error("assign outputs before returning them")); + ORT_RETURN_IF(pattern_.Nodes().empty(), Error("return must identify the final operator outputs")); + const auto& outputs = pattern_.Nodes().back().outputs; + ORT_RETURN_IF(names.size() != outputs.size(), Error("return must preserve all final outputs in order")); + 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(found->second.handle != outputs[i].handle, + Error("return must preserve all final outputs in order")); + } + returned_values = true; + } else { + for (const auto& name : names) ORT_RETURN_IF(symbols_.count(name), Error("duplicate assignment name")); + ORT_RETURN_IF_ERROR(Expect("=")); + } } } - ORT_RETURN_IF_ERROR(Call(names)); + if (!returned_values) ORT_RETURN_IF_ERROR(Call(names)); if (returned) { 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; } @@ -530,7 +548,7 @@ class PatternParser { (c == ':' && position_ < text_.size() && text_[position_] == ':')) { ++position_; } else { - ORT_RETURN_IF(std::string_view("(),=;").find(c) == std::string_view::npos, Error("unexpected character")); + 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(); diff --git a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc index bdf691a325a37..4a7049fb0c76c 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc @@ -48,11 +48,12 @@ Rules MakeExampleRules() { Rules rules; { Pattern p(R"( - (x) => + (x) => { add = Add(x, 3.0) clip = Clip(add, 0.0, 6.0) mul = Mul(clip, x) return Div(mul, 6.0) + } )"); auto guard = [](const Graph&, const Match& m) { return IsFloat32(*m.valueOf("x")); }; Replacement replacement("(x) => HardSwish(x)"); @@ -61,9 +62,10 @@ Rules MakeExampleRules() { } { Pattern p(R"( - (x) => + (x) => { sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) return Mul(sigmoid, x) + } )"); auto guard = [](const Graph&, const Match& m) { return IsFloat32(*m.valueOf("x")); }; Replacement replacement("(x) => HardSwish(x)"); @@ -72,7 +74,7 @@ Rules MakeExampleRules() { } #ifndef DISABLE_CONTRIB_OPS { - Pattern p("(x, y, divisor) => mm = MatMul(x, y); return Div(mm, divisor)"); + Pattern p("(x, y, divisor) => { mm = MatMul(x, y); return Div(mm, divisor); }"); auto guard = [](const Graph& graph, const Match& m) { auto scalar = Scalar(graph, *m.valueOf("divisor")); return IsFloat32(*m.valueOf("x")) && IsFloat32(*m.valueOf("y")) && @@ -87,7 +89,7 @@ Rules MakeExampleRules() { rules.push_back(std::move(rule)); } { - Pattern p("(x, y) => transpose = Transpose(y); return MatMul(x, transpose)"); + Pattern p("(x, y) => { transpose = Transpose(y); return MatMul(x, transpose); }"); auto guard = [](const Graph&, const Match& m) { if (!IsFloat32(*m.valueOf("x")) || !IsFloat32(*m.valueOf("y"))) return false; const auto* shape = m.valueOf("y")->Shape(); @@ -110,7 +112,7 @@ Rules MakeExampleRules() { rules.push_back(std::move(rule)); } { - Pattern p("(x, y, divisor) => mm = com.microsoft::FusedMatMul(x, y); return Div(mm, divisor)"); + Pattern p("(x, y, divisor) => { mm = com.microsoft::FusedMatMul(x, y); return Div(mm, divisor); }"); auto guard = [](const Graph& graph, const Match& m) { auto scalar = Scalar(graph, *m.valueOf("divisor")); auto it = m.nodes[0]->GetAttributes().find("alpha"); @@ -141,7 +143,7 @@ Rules MakeExampleRules() { } { Pattern p(R"( - (x, scale, epsilon) => + (x, scale, epsilon) => { square = Pow(x, 2.0) mean = ReduceMean(square) add = Add(mean, epsilon) @@ -149,6 +151,7 @@ Rules MakeExampleRules() { reciprocal = Reciprocal(sqrt) normalize = Mul(x, reciprocal) return Mul(normalize, scale) + } )"); auto guard = [](const Graph& graph, const Match& m) { if (!IsFloat32(*m.valueOf("x")) || !IsFloat32(*m.valueOf("scale"))) return false; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index e27ef512f0905..4e405e302dfd9 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -83,9 +83,9 @@ Rule IdentityPair(const char* name = "identity", bool accept = true) { } Rule IdentitySplit() { - Pattern p("(x) => first = Identity(x); return a, b = Split(first)"); + Pattern p("(x) => { first = Identity(x); a, b = Split(first); return a, b; }"); auto guard = [](const Graph&, const Match&) { return true; }; - Replacement plan("(x, axis) => return a, b = Split(x, axis=axis)"); + Replacement plan("(x, axis) => { a, b = Split(x, axis=axis); return a, b; }"); auto replacement = [plan](const Graph&, const Match& match) { const auto& node = *match.nodes.back(); auto axis = node.GetAttributes().find("axis"); @@ -360,7 +360,7 @@ TEST(PatternRewriterTests, NamedMatchLookupAndReplacementSafety) { test.builder.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); auto first = IdentityPair(), second = IdentityPair(); - second.pattern = Pattern("(input) => middle = Identity(input); return root = Identity(middle)"); + second.pattern = Pattern("(input) => { middle = Identity(input); root = Identity(middle); return root; }"); first.guard = [=](const Graph&, const Match& m) { EXPECT_EQ(m.valueOf("x"), x); EXPECT_EQ(m.valueOf("a"), intermediate); @@ -484,8 +484,8 @@ TEST(PatternRewriterTests, ReplacementBranchesAndMultiOutputRoot) { std::make_unique(Rules{branched})); auto split = IdentitySplit(); - split.pattern = Pattern("(x) => a = Identity(x); b = Identity(a); return left, right = Split(b)"); - split.replacement = Replacement("(x) => a=Identity(x); return left, right = Split(a, axis=0)"); + 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(); @@ -1424,12 +1424,46 @@ TEST(PatternRewriterTests, ParserOnlyConstruction) { 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);return result=Identity(a)")); - EXPECT_NO_THROW(Replacement("(x)=>return result=Identity(x)")); + 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;}"}) { + SCOPED_TRACE(text); + EXPECT_NO_THROW(Pattern(text)); + EXPECT_NO_THROW(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(a);return a;}", + "(x)=>{left,right=Split(x);return left;}", "(x)=>{left,right=Split(x);return right,left;}", + "(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(Pattern(text), OnnxRuntimeException); + EXPECT_THROW(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, ParsedMultiOutputDefinitionsAndLimits) { - Pattern p("(x)=>a=Identity(x);b=ai.onnx::Identity(a);return left,right = ai.onnx::Split(b)"); + 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); @@ -1440,10 +1474,10 @@ TEST(PatternRewriterTests, ParsedMultiOutputDefinitionsAndLimits) { 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)=>return a,b = Split(missing)"), OnnxRuntimeException); - EXPECT_THROW(p = Pattern("(x)=>return left,left = Split(x)"), OnnxRuntimeException); - EXPECT_THROW(p = Pattern("(x)=>return x,right = Split(x)"), OnnxRuntimeException); - EXPECT_THROW(p = Pattern("(x)=>a,b=Split(x);return a,b"), 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; @@ -1453,14 +1487,16 @@ TEST(PatternRewriterTests, ParsedMultiOutputDefinitionsAndLimits) { } return text; }; - EXPECT_THROW(Pattern("(x)=>return " + outputs_text(0, 65) + "=Split(x)"), OnnxRuntimeException); + 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 + "return " + outputs_text(3, 64) + "=Split(v2_0)"), OnnxRuntimeException); + EXPECT_THROW(Pattern(many_values + outputs_text(3, 64) + "=Split(v2_0);return " + outputs_text(3, 64)), + OnnxRuntimeException); } TEST(PatternRewriterTests, MultiOutputIntermediateBoundariesAndBindings) { @@ -1528,7 +1564,8 @@ TEST(PatternRewriterTests, MultiOutputRootEdgesAndInvalidRecipe) { test.builder.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); Statistics stats; - for (auto text : {"(x)=>Identity(x)", "(x)=>Split(x)", "(x)=>return left,middle,right=Split(x)"}) { + 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()); From af63ee2920f80766fb0db0f9c19c91e8ceffe8df Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 13 Oct 2026 22:29:06 -0700 Subject: [PATCH 08/45] Disambiguate parser construction expressions in syntax tests Use explicit void-cast brace expressions in Google Test exception macros to avoid declaration parsing and default-constructor lookup. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/optimizer/pattern_rewriter_test.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 4e405e302dfd9..af43932cf211f 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1436,8 +1436,8 @@ TEST(PatternRewriterTests, BracedBodiesAndNamedReturns) { "(x)=>left,right=Split(x);return left,right", "(x)=>{left,right=ai.onnx::Split(x);return left,right;}"}) { SCOPED_TRACE(text); - EXPECT_NO_THROW(Pattern(text)); - EXPECT_NO_THROW(Replacement(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)")); @@ -1452,8 +1452,8 @@ TEST(PatternRewriterTests, BracedBodiesAndNamedReturns) { "(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(Pattern(text), OnnxRuntimeException); - EXPECT_THROW(Replacement(text), OnnxRuntimeException); + 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;}"); From e06a09fcd72543d9cd3ef9e3b63d97fc5a2e1bb9 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 14 Oct 2026 07:56:11 -0700 Subject: [PATCH 09/45] Support dotted pattern operators and rank-one literal lists Parse dotted domain names in patterns as well as replacements. Match and synthesize numeric rank-one tensor literals in positional inputs and exact numeric list attributes in keyword arguments, with parsing, structural, rollback and numerical coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 46 ++++- .../core/optimizer/pattern_rewriter.cc | 140 ++++++++++++--- onnxruntime/core/optimizer/pattern_rewriter.h | 7 +- .../test/optimizer/pattern_rewriter_test.cc | 168 +++++++++++++++++- 4 files changed, 328 insertions(+), 33 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 74d21f5174145..2810a96dff53b 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -81,16 +81,18 @@ A final call without an assignment, such as `return Div(product, 6.0)`, also des but does not give that output a lookup name. Replacements additionally accept a final call without the `return` keyword, as above. -Use `domain::Operator(...)` for qualified operators: +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)"); +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` -is equivalent to `Identity`. Replacements also accept -`com.microsoft.FusedMatMul(...)`; patterns use the `::` form. +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 @@ -109,6 +111,21 @@ 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 @@ -121,13 +138,24 @@ 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. +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, -not attribute lists or custom predicate syntax. Inspect `match.nodes` in the C++ -guard for ranges, list-valued attributes, shapes, or other semantic conditions. +and numeric list equality, but not custom predicate syntax or string lists. +Inspect `match.nodes` in the C++ guard for ranges, shapes, or other semantic conditions. Positional inputs must precede keyword attributes. -Numeric literals in replacements synthesize rank-zero float32 or int64 -initializers. Replacement keyword literals set attributes on the new node. +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 diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index eae1ef96f710c..5dccd70aae158 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -8,7 +8,9 @@ #include #include #include +#include #include +#include #include "core/graph/graph_utils.h" #include "core/optimizer/initializer.h" @@ -28,7 +30,17 @@ std::string Signature(const Pattern& pattern) { auto it = pattern.Constants().find(id); if (it == pattern.Constants().end()) continue; stream << 'c' << id << ':' << it->second.index() << ':'; - std::visit([&](auto value) { stream << value; }, it->second); + 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()) { @@ -98,18 +110,28 @@ bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, } for (const auto& [handle, expected] : pattern.Constants()) { const auto* tensor = graph_utils::GetConstantInitializer(graph, match.values[handle]->Name()); - if (!tensor || tensor->dims_size() != 0 || + const bool vector = expected.index() >= 2; + if (!tensor || tensor->dims_size() != (vector ? 1 : 0) || tensor->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL) return false; - const bool floating = std::holds_alternative(expected); + 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 false; Initializer initializer(graph, *tensor, graph.ModelPath()); - if (initializer.size() != 1) return false; - if (floating ? *initializer.data() != std::get(expected) - : *initializer.data() != std::get(expected)) - return false; + 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 false; } for (size_t i = 0; i + 1 < match.nodes.size(); ++i) { const Node& node = *match.nodes[i]; @@ -197,16 +219,32 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R if (const auto* scalar = std::get_if(&nodes[i].inputs[j])) { ONNX_NAMESPACE::TensorProto tensor; tensor.set_name(graph.GenerateNodeArgName("rewrite_literal")); - if (const auto* floating = std::get_if(&scalar->value)) { - tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - tensor.add_float_data(*floating); - } else { - tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); - tensor.add_int64_data(std::get(scalar->value)); - } + 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()); - type.mutable_tensor_type()->mutable_shape(); + 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); @@ -341,6 +379,12 @@ Value Pattern::Constant(int64_t 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 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); @@ -548,7 +592,7 @@ class PatternParser { (c == ':' && position_ < text_.size() && text_[position_] == ':')) { ++position_; } else { - ORT_RETURN_IF(std::string_view("(),=;{}").find(c) == std::string_view::npos, Error("unexpected character")); + 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(); @@ -586,6 +630,33 @@ class PatternParser { 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) { ORT_RETURN_IF(kind_ != Kind::Identifier || token_ == "return", Error("expected operator name")); Operator op{token_, ""}; @@ -596,7 +667,7 @@ class PatternParser { ORT_RETURN_IF(kind_ != Kind::Identifier, Error("expected qualified operator name")); op.name = token_; ORT_RETURN_IF_ERROR(Next()); - } else if (replacement_attributes_) { + } else { auto dot = op.name.rfind('.'); if (dot != std::string::npos) { op.domain = op.name.substr(0, dot); @@ -634,6 +705,29 @@ class PatternParser { 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)); @@ -657,10 +751,16 @@ class PatternParser { } } else { ORT_RETURN_IF(keyword, Error("positional input after attribute")); - std::variant number; - ORT_RETURN_IF_ERROR(Number(number)); ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("too many values")); - inputs.push_back(std::visit([&](auto value) { return pattern_.Constant(value); }, number)); + 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; diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index a01419fb7e31a..371374c08372e 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -46,6 +46,8 @@ struct Operator { std::string domain; }; +using Literal = std::variant, InlinedVector>; + struct OpsetRequirement { std::string domain; int min_version; @@ -78,6 +80,7 @@ class Pattern { 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); @@ -85,7 +88,7 @@ class Pattern { size_t value_count_ = 0; InlinedVector nodes_; InlinedVector values_; - InlinedHashMap> constants_; + InlinedHashMap constants_; Value AddValue(std::string_view name); }; @@ -108,7 +111,7 @@ struct Replacement { size_t slot = 0; }; struct Scalar { - std::variant value; + Literal value; }; using Input = std::variant; struct NodeRecipe { diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index af43932cf211f..b21aa33bdf2c6 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -505,7 +505,7 @@ TEST(PatternRewriterTests, ReplacementBranchesAndMultiOutputRoot) { 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)=>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)", @@ -1004,7 +1004,7 @@ TEST(PatternRewriterTests, ParsedAttributesMultiOutputsAndDomains) { "(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(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 { @@ -1055,6 +1055,169 @@ TEST(PatternRewriterTests, LiteralTypeAndRankMatching) { 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}) { + for (int scenario = 0; scenario < 5; ++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,[1,3]);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; @@ -2020,6 +2183,7 @@ TEST(PatternRewriterTests, NumericalFusedMatMulDivPreservesAttributes) { } }; 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})); } From ec7c6b71496d863d9637cedbf184ca87edb5756b Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 14 Oct 2026 08:19:53 -0700 Subject: [PATCH 10/45] Keep rejected vector rule disjoint from all negative test inputs The mismatching [1,3] graph correctly matched the additional [1,3] rule. Use [4,5] for the always-rejected structural-plan regression and trace mode/scenario explicitly. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/optimizer/pattern_rewriter_test.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index b21aa33bdf2c6..8ef41391dc68a 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1080,7 +1080,9 @@ TEST(PatternRewriterTests, VectorLiteralParsingAndErrors) { 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); @@ -1105,7 +1107,7 @@ TEST(PatternRewriterTests, VectorLiteralMatchingAndNumerics) { 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,[1,3]);return Identity(a);}"); + 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); From 8d8c7cb3f66b4e49c2057e113d944f953fea712b Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 14 Oct 2026 09:55:10 -0700 Subject: [PATCH 11/45] Extend rewriter boundaries, attribute captures, optional inputs and diagnostics Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 124 ++++- .../core/optimizer/pattern_rewriter.cc | 301 +++++++++--- onnxruntime/core/optimizer/pattern_rewriter.h | 42 +- .../test/optimizer/pattern_rewriter_test.cc | 435 +++++++++++++++++- 4 files changed, 810 insertions(+), 92 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 2810a96dff53b..fdf0a7568a3d8 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -57,8 +57,8 @@ change evaluation order; validate equivalence within the intended tolerance. 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 all -outputs of the final assignment. Whitespace separates calls, and semicolons are +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. Nothing may follow the closing brace. Comments are not supported inside definitions. @@ -102,6 +102,27 @@ values. Unknown names and unbound values raise explicit errors. Read-only `Pattern::NamedValue`, `Nodes`, `Constants`, `ValueCount`, and `IsInput` are available for introspection. +### 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 `?` to allow either a present value or an absent slot: + +```cpp +Pattern pattern("(x,low?,high?) => { 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?) => Clip(x,v,v)` does not match +one present and one absent slot. `match.valueOf("low")` returns `nullptr` when +the optional binding is absent; guards 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** @@ -153,6 +174,23 @@ and numeric list equality, but not custom predicate syntax or string lists. Inspect `match.nodes` in the C++ guard 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 guards and replacements without a callback. +Capture names are separate from tensor values and cannot collide with them. +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`. Missing attributes +are not captured from schema defaults: only the resolved node's map is used. + 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. @@ -194,17 +232,33 @@ Pattern pattern(R"( Replacement replacement("(x) => { left,right = Split(x,axis=0); return left,right; }"); ``` -This replacement preserves the root's two original output `NodeArg`s in order, +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 only one output and cannot replace a two-output root. -A named return must list **all outputs of the final defined operator in their -original order**. Returning a leaf input, an earlier operator's outputs, a subset, -duplicate names, or reordered outputs is rejected. A named return does not create -an additional operator or introduce output selection. +`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: -Intermediate operators may have unused outputs, but **none** of their outputs -may be exposed outside the matched deletion set. +```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. ## Computed attributes with a replacement callback @@ -260,8 +314,8 @@ Use local bindings, attributes, and constants to establish applicability. `update` accepts a concrete `AttributeProto`, including list-valued attributes. It cannot overwrite a named pattern value or store a referenced/undefined -attribute. Extra bindings are cleared before each replacement callback; -`attributeOf` retrieves one and errors if it is missing. A replacement header +attribute. Extra bindings are cleared before each replacement callback; captured attributes +remain available. `attributeOf` retrieves either kind 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. @@ -275,19 +329,20 @@ must be connected to their final operator. Before calling a rule's guard, the matcher checks: -- Exact operator name/domain, ordered input/output counts, and repeated binding - identity. Missing optional slots and variadic patterns are not supported. +- 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 non-root output stays internal. A graph output, an external consumer, or - an implicit capture by an external subgraph prevents the rewrite. +- 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. -The root may have graph outputs or external consumers, because its complete -output list is preserved. Even match-only mode applies deletion-boundary checks: +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 @@ -375,7 +430,7 @@ mode with the same eligibility and priority semantics. Every pattern must contain at least two nodes. A replacement must have a nonempty node sequence with **strictly fewer nodes** than the pattern, and its -final output count must equal the root's output count. Expansion and one-node +returned-value count must equal the pattern's returned-value count. Expansion and one-node rewrites are rejected. Each successful rewrite reduces live node count, so a graph with N nodes admits at most N-1 rewrites. This establishes termination, not a time bound for arbitrary callbacks or graph fanout. @@ -404,11 +459,40 @@ candidates, structural checks, guard calls, successful matches, and completed rewrites. Match-only counts all passing guards; 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, constants, attributes, deletion +boundaries, execution providers, guards, 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 guard 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 records or reason strings are allocated when +diagnostics are disabled. 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. + ## Limits and unsupported forms Definitions are limited to 64 nodes, 256 values, and 64 inputs/outputs per operator. They are connected acyclic graphs rooted at the final operator. -There is no alternation, implicit commutation, optional/variadic matching, +There is no 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`. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 5dccd70aae158..6d67535c86197 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -25,6 +25,9 @@ std::string Key(const std::string& op, const std::string& domain) { return domai std::string Signature(const Pattern& pattern) { std::ostringstream stream; 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); @@ -57,15 +60,21 @@ std::string Signature(const Pattern& pattern) { using Producers = InlinedHashMap; bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, const Producers& producers, - Match& match) { + Match& match, FailureStage& stage, std::string_view& 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 false; + 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]; @@ -78,47 +87,67 @@ bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, break; } } - if (!node || !distinct.insert(node->Index()).second || node->OpType() != expected.op.name || - node->Domain() != expected.op.domain || node->InputDefs().size() != expected.inputs.size() || - node->OutputDefs().size() != expected.outputs.size() || !node->GetAttributeNameToSubgraphMap().empty() || - !node->ImplicitInputDefs().empty() || node->GetLayeringAnnotation() != root.GetLayeringAnnotation()) { - return false; - } + 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() || !attribute.Matches(found->second)) return false; + if (found == node->GetAttributes().end() || !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 false; + 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 = node->InputDefs()[j]; - const NodeArg*& binding = match.values[expected.inputs[j]]; - if (!actual->Exists() || (binding && binding != actual)) return false; + 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. 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 false; + 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 false; + 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 false; + 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; @@ -131,13 +160,20 @@ bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, } }, expected); - if (!equal) return false; + if (!equal) return fail(FailureStage::Constant, "Literal initializer values mismatch."); } - for (size_t i = 0; i + 1 < match.nodes.size(); ++i) { - const Node& node = *match.nodes[i]; - if (graph.NodeProducesGraphOutput(node)) return false; + 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())) return false; + 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; @@ -156,8 +192,16 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R ORT_RETURN_IF(nodes.empty() || nodes.size() >= match.nodes.size(), "Replacements must contain fewer nodes than the matched pattern."); const Node& root = *match.nodes.back(); - ORT_RETURN_IF(nodes.back().output_count != root.OutputDefs().size(), - "The final replacement must preserve the root output count."); + ORT_RETURN_IF(recipe.Outputs().size() != pattern.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] = graph.GetNodeArg(match.values[pattern.Outputs()[i]]->Name()); + } InlinedVector> external_inputs; for (size_t i = 0; i < nodes.size(); ++i) { const auto& node = nodes[i]; @@ -172,6 +216,10 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R "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."); @@ -193,11 +241,12 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R ORT_RETURN_IF( value.handle < 0 || static_cast(value.handle) >= match.values.size() || !pattern.IsInput(value), "Replacement inputs must be pattern input variables."); - inputs.push_back(graph.GetNodeArg(match.values[value]->Name())); + const NodeArg* binding = match.values[value]; + inputs.push_back(binding ? graph.GetNodeArg(binding->Name()) : &graph.GetOrCreateNodeArg("", nullptr)); } external_inputs.push_back(std::move(inputs)); } - // Require a connected sequence ending at the node that inherits the root outputs. + // 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."); @@ -251,12 +300,9 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R } } InlinedVector outputs; - if (i + 1 == nodes.size()) { - for (const NodeArg* output : root.OutputDefs()) outputs.push_back(graph.GetNodeArg(output->Name())); - } else { - for (size_t j = 0; j < nodes[i].output_count; ++j) - outputs.push_back(&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("rewrite_value"), nullptr)); - } + 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; NodeAttributes resolved_attributes; @@ -288,20 +334,68 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R } replacement.SetExecutionProviderType(root.GetExecutionProviderType()); } + // 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 source_arg; int destination_arg; }; InlinedVector outgoing; - for (auto it = root.OutputEdgesBegin(); it != root.OutputEdgesEnd(); ++it) { - outgoing.push_back({it->GetNode().Index(), it->GetSrcArgIndex(), it->GetDstArgIndex()}); + 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())); @@ -314,6 +408,7 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R 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; @@ -327,8 +422,12 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R graph.UpdateProducerNode(output->Name(), index); } } - for (const auto& edge : outgoing) - graph.AddEdge(added.back(), edge.destination, edge.source_arg, edge.destination_arg); + 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(); } @@ -426,12 +525,12 @@ InlinedVector Pattern::MultiOp(Operator op, gsl::span inputs "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 >= 0 && static_cast(input.handle) < value_count_, + 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(values_[input.handle]); + 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; } @@ -459,7 +558,12 @@ class PatternParser { ORT_RETURN_IF_ERROR(Identifier(name)); ORT_RETURN_IF(symbols_.count(name), Error("duplicate input name")); ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("too many values")); - symbols_.emplace(name, pattern_.Input(name)); + Value input = pattern_.Input(name); + symbols_.emplace(name, input); + if (token_ == "?") { + pattern_.optional_inputs_.insert(input.handle); + ORT_RETURN_IF_ERROR(Next()); + } if (token_ != ",") break; ORT_RETURN_IF_ERROR(Next()); } while (true); @@ -502,24 +606,25 @@ class PatternParser { Error("expected value identifier")); if (returned) { ORT_RETURN_IF(token_ == "=", Error("assign outputs before returning them")); - ORT_RETURN_IF(pattern_.Nodes().empty(), Error("return must identify the final operator outputs")); - const auto& outputs = pattern_.Nodes().back().outputs; - ORT_RETURN_IF(names.size() != outputs.size(), Error("return must preserve all final outputs in order")); + ORT_RETURN_IF(pattern_.Nodes().empty(), 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(found->second.handle != outputs[i].handle, - Error("return must preserve all final outputs in order")); + ORT_RETURN_IF(pattern_.IsInput(found->second), Error("cannot return a leaf input")); + pattern_.outputs_.push_back(found->second); } returned_values = true; } else { - for (const auto& name : names) ORT_RETURN_IF(symbols_.count(name), Error("duplicate assignment name")); + for (const auto& name : names) + ORT_RETURN_IF(symbols_.count(name) || pattern_.attribute_names_.count(name), + Error("duplicate assignment or capture name")); ORT_RETURN_IF_ERROR(Expect("=")); } } } if (!returned_values) ORT_RETURN_IF_ERROR(Call(names)); if (returned) { + if (!returned_values) pattern_.outputs_ = pattern_.Nodes().back().outputs; 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")); @@ -535,7 +640,8 @@ class PatternParser { [&](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) reachable.insert(input.handle); + for (Value input : it->inputs) + if (input.handle >= 0) reachable.insert(input.handle); } ORT_RETURN_IF(reachable.size() != pattern_.ValueCount(), Error("unused input")); destination = std::move(pattern_); @@ -592,7 +698,7 @@ class PatternParser { (c == ':' && position_ < text_.size() && text_[position_] == ':')) { ++position_; } else { - ORT_RETURN_IF(std::string_view("(),=;{}[]").find(c) == std::string_view::npos, Error("unexpected character")); + 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(); @@ -604,7 +710,8 @@ class PatternParser { } Status Identifier(std::string& name, bool allow_domain = false) { ORT_RETURN_IF( - kind_ != Kind::Identifier || token_ == "return" || (!allow_domain && token_.find('.') != std::string::npos), + kind_ != Kind::Identifier || token_ == "return" || token_ == "None" || + (!allow_domain && token_.find('.') != std::string::npos), Error("expected value identifier")); name = token_; return Next(); @@ -689,15 +796,23 @@ class PatternParser { 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 (replacement_bindings_ && kind_ == Kind::Identifier) { - auto found = symbols_.find(token_); - ORT_RETURN_IF(found == symbols_.end() || !pattern_.IsInput(found->second), - Error("attribute binding must reference a declared replacement input")); - bindings.emplace(name, token_); - attribute_inputs_.insert(found->second.handle); - ORT_RETURN_IF_ERROR(Next()); + if (kind_ == Kind::Identifier) { + std::string capture; + ORT_RETURN_IF_ERROR(Identifier(capture)); + 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 { + ORT_RETURN_IF(symbols_.count(capture), Error("attribute capture conflicts with a value name")); + pattern_.attribute_names_.insert(capture); + attributes.push_back({name, [](const auto&) { return true; }, capture}); + } } else if (kind_ == Kind::String) { attributes.push_back(AttributeMatcher::Equal(name, token_)); auto& attr = literal_attributes[name]; @@ -745,9 +860,13 @@ class PatternParser { } } else { ORT_RETURN_IF(keyword, Error("positional input after attribute")); - auto found = symbols_.find(name); - ORT_RETURN_IF(found == symbols_.end(), Error("undefined input value")); - inputs.push_back(found->second); + 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")); @@ -819,6 +938,10 @@ Replacement::Replacement(std::string_view text) { 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}); @@ -837,6 +960,7 @@ Replacement::Replacement(std::string_view text) { 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)); } @@ -845,7 +969,8 @@ const NodeArg* Match::valueOf(std::string_view name) const { auto value = std::find_if(pattern->values_.begin(), pattern->values_.end(), [&](const Value& candidate) { return candidate.name == name; }); ORT_ENFORCE(!name.empty() && value != pattern->values_.end(), "Unknown pattern value: ", name); - ORT_ENFORCE(value->handle >= 0 && static_cast(value->handle) < values.size() && values[value->handle], + 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]; } @@ -858,17 +983,30 @@ void Match::update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) 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)), + "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 { + auto captured = captured_attributes_.find(std::string(name)); + if (captured != captured_attributes_.end()) return captured->second; 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 (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(); +} + PatternRewriter::PatternRewriter(Rules rules, Options options) : GraphTransformer("PatternRewriter"), rules_(std::move(rules)), options_(options) { InlinedHashMap signatures; @@ -917,7 +1055,8 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) 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) reachable.insert(input.handle); + for (Value input : it->inputs) + if (input.handle >= 0) reachable.insert(input.handle); } ORT_ENFORCE(reachable.size() == rule.pattern.ValueCount(), "Unused pattern input."); invalidation_depth_ = std::max(invalidation_depth_, rule.pattern.Nodes().size()); @@ -940,6 +1079,10 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) 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; @@ -950,7 +1093,10 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { imported->second <= requirement.max_version; }); applicable.push_back(supported); - if (supported) applicable_rules.push_back(id); + 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; @@ -986,13 +1132,27 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { statistics.candidates += rule_ids.size(); ++statistics.structural_checks; Match match; - if (!MatchPattern(pattern, graph, *root, producers, match) || !CanRewritePartition(match, options_)) return; + FailureStage stage = FailureStage::Structure; + std::string_view reason; + if (!MatchPattern(pattern, graph, *root, producers, match, stage, reason)) { + for (size_t id : rule_ids) report(id, root->Index(), stage, reason); + return; + } + if (!CanRewritePartition(match, options_)) { + for (size_t id : rule_ids) + report(id, root->Index(), FailureStage::ExecutionProvider, "Mixed or unsupported execution providers."); + return; + } for (size_t id : rule_ids) { if (!options_.match_only && id >= winner) break; match.pattern = &rules_[id].pattern; match.ClearUpdates(); ++statistics.guards; - if (!rules_[id].guard(graph, match)) continue; + if (!rules_[id].guard(graph, match)) { + report(id, root->Index(), FailureStage::Guard, + match.RejectionReason().empty() ? std::string_view("Guard rejected match.") : match.RejectionReason()); + continue; + } if (options_.match_only) ++statistics.matches; if (id < winner) { winner = id; @@ -1029,15 +1189,22 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { InlinedVector added; winning_match.ClearUpdates(); auto recipe = rules_[winner].replacement(graph, winning_match); - ORT_RETURN_IF(!recipe.Plan(), "Replacement has no parsed plan."); 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; }); }; - for (const auto& node : *recipe.Plan()) - ORT_RETURN_IF(!declared(node.domain), "Replacement domain has no rule opset requirement: ", node.domain); - ORT_RETURN_IF_ERROR(Replace(graph, *winning_pattern, winning_match, recipe, producers, added)); + 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. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 371374c08372e..692a6bd9994c1 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -31,6 +31,7 @@ struct Value { struct AttributeMatcher { std::string name; std::function predicate; + std::string capture; 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}); } @@ -70,6 +71,9 @@ class Pattern { 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_; } private: friend class detail::PatternParser; @@ -89,6 +93,9 @@ class Pattern { InlinedVector nodes_; InlinedVector values_; InlinedHashMap constants_; + InlinedVector outputs_; + InlinedHashSet optional_inputs_; + InlinedHashSet attribute_names_; Value AddValue(std::string_view name); }; @@ -98,11 +105,22 @@ struct Match { const NodeArg* valueOf(std::string_view name) const; void update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) const; const ONNX_NAMESPACE::AttributeProto& attributeOf(std::string_view name) const; - void ClearUpdates() const { extra_attributes_.clear(); } + void ClearUpdates() const { + extra_attributes_.clear(); + rejection_reason_.clear(); + } + bool BindAttribute(std::string_view name, const ONNX_NAMESPACE::AttributeProto& value); + bool Reject(std::string_view reason) const { + rejection_reason_ = reason; + return false; + } + const std::string& RejectionReason() const { return rejection_reason_; } const Pattern* pattern = nullptr; private: + NodeAttributes captured_attributes_; mutable NodeAttributes extra_attributes_; + mutable std::string rejection_reason_; }; struct Replacement { @@ -113,7 +131,7 @@ struct Replacement { struct Scalar { Literal value; }; - using Input = std::variant; + using Input = std::variant; struct NodeRecipe { std::string op; std::string domain; @@ -126,10 +144,12 @@ struct Replacement { // 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_; } private: friend class PatternRewriter; std::shared_ptr> plan_; + InlinedVector outputs_; }; struct Rule { @@ -153,11 +173,29 @@ struct Statistics { size_t rewrites = 0; }; +enum class FailureStage { Opset, + Structure, + Constant, + Attribute, + Boundary, + ExecutionProvider, + Guard, + 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; }; // Immutable compiled rules; statistics belong to a call, not this transformer. diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 8ef41391dc68a..bc3585aadf074 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1599,7 +1599,11 @@ TEST(PatternRewriterTests, BracedBodiesAndNamedReturns) { "(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)=>{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}); @@ -1610,8 +1614,8 @@ TEST(PatternRewriterTests, BracedBodiesAndNamedReturns) { "(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(a);return a;}", - "(x)=>{left,right=Split(x);return left;}", "(x)=>{left,right=Split(x);return right,left;}", + "(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)", @@ -1627,6 +1631,431 @@ TEST(PatternRewriterTests, BracedBodiesAndNamedReturns) { 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.guard = [&guards](const Graph&, const 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?,high?)=>{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}); + 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?,high?)=>{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.guard = [scenario, &guards](const Graph&, const 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 guard must not retain references to a previous graph's match. + rule.guard = [](const Graph&, const 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.guard = [&guards](const Graph&, const 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.guard = [](const Graph&, const 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?)=>{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.guard = [](const Graph&, const 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::Guard + : 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, 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"); From 8cd0964692315d462bb08fc37a1849c36c727c23 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 14 Oct 2026 09:56:52 -0700 Subject: [PATCH 12/45] Clarify rewriter introspection and diagnostic allocation semantics Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index fdf0a7568a3d8..206e8c9b8ec3c 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -99,7 +99,8 @@ 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`, `ValueCount`, and `IsInput` +`Pattern::NamedValue`, `Nodes`, `Constants`, `Outputs`, `ValueCount`, `IsInput`, +`IsOptional`, and `AttributeNames` are available for introspection. ### Optional input slots @@ -483,8 +484,9 @@ A guard 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 records or reason strings are allocated when -diagnostics are disabled. Replacement validation status errors are reported and +records for later inspection. No diagnostic records are constructed when +diagnostics are disabled; `Reject` still stores an explicitly supplied guard +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. From be3e0d165de342e681a7dc6dd0734da11703a8e1 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 14 Oct 2026 10:35:02 -0700 Subject: [PATCH 13/45] Allow one-node and nonshrinking rewrites with bounded candidate iterations Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 38 +++- .../core/optimizer/pattern_rewriter.cc | 209 +++++++++--------- onnxruntime/core/optimizer/pattern_rewriter.h | 3 + .../test/optimizer/pattern_rewriter_test.cc | 111 +++++++++- 4 files changed, 250 insertions(+), 111 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 206e8c9b8ec3c..066be14d3d74f 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -1,7 +1,7 @@ # Writing C++ graph rewrite rules The experimental `onnxruntime::rewrite::PatternRewriter` replaces connected -operator patterns with smaller, semantically equivalent graphs. Include +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. Guards and replacement callbacks remain ordinary C++ functions. @@ -429,17 +429,39 @@ passing guard. With `Options{.match_only=true}`, all passing guards are counted and replacements are not called. `indexed=false` selects the reference matching mode with the same eligibility and priority semantics. -Every pattern must contain at least two nodes. A replacement must have a -nonempty node sequence with **strictly fewer nodes** than the pattern, and its -returned-value count must equal the pattern's returned-value count. Expansion and one-node -rewrites are rejected. Each successful rewrite reduces live node count, so a -graph with N nodes admits at most N-1 rewrites. This establishes termination, -not a time bound for arbitrary callbacks or graph fanout. +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 without retaining pointers to deleted nodes. +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 diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 6d67535c86197..0402308c0bba6 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -189,8 +189,7 @@ bool CanRewritePartition(const Match& match, const Options& options) { Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const Replacement& recipe, Producers& producers, InlinedVector& added) { const auto& nodes = *recipe.Plan(); - ORT_RETURN_IF(nodes.empty() || nodes.size() >= match.nodes.size(), - "Replacements must contain fewer nodes than the matched pattern."); + ORT_RETURN_IF(nodes.empty(), "Replacements must contain at least one node."); const Node& root = *match.nodes.back(); ORT_RETURN_IF(recipe.Outputs().size() != pattern.Outputs().size(), "Replacement return count must equal the pattern return count."); @@ -1009,11 +1008,12 @@ bool Match::BindAttribute(std::string_view name, const ONNX_NAMESPACE::Attribute PatternRewriter::PatternRewriter(Rules rules, Options options) : GraphTransformer("PatternRewriter"), rules_(std::move(rules)), options_(options) { + 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().size() >= 2 && rule.guard && rule.replacement, - "Rules need a name, at least two nodes, a guard and a recipe."); + ORT_ENFORCE(!rule.name.empty() && !rule.pattern.Nodes().empty() && rule.guard && rule.replacement, + "Rules need a name, at least one node, a guard and a recipe."); InlinedHashSet domains; for (auto& requirement : rule.opsets) { if (requirement.domain == kOnnxDomainAlias) requirement.domain = kOnnxDomain; @@ -1026,8 +1026,6 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) } if (auto* recipe = rule.replacement.target(); recipe && recipe->plan_) { auto nodes = *recipe->plan_; - ORT_ENFORCE(nodes.size() < rule.pattern.Nodes().size(), - "Replacements must contain fewer nodes than the matched pattern."); 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); @@ -1119,106 +1117,115 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { if (output->Exists()) producers[output] = index; } } - // Nodes are revisited only after a strictly node-count-decreasing rewrite. - for (size_t cursor = 0; cursor < worklist.size(); ++cursor) { - queued.erase(worklist[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; - auto consider = [&](const Pattern& pattern, gsl::span rule_ids) { - statistics.candidates += rule_ids.size(); - ++statistics.structural_checks; - Match match; - FailureStage stage = FailureStage::Structure; - std::string_view reason; - if (!MatchPattern(pattern, graph, *root, producers, match, stage, reason)) { - for (size_t id : rule_ids) report(id, root->Index(), stage, reason); - return; - } - if (!CanRewritePartition(match, options_)) { - for (size_t id : rule_ids) - report(id, root->Index(), FailureStage::ExecutionProvider, "Mixed or unsupported execution providers."); - return; - } - for (size_t id : rule_ids) { - if (!options_.match_only && id >= winner) break; - match.pattern = &rules_[id].pattern; - match.ClearUpdates(); - ++statistics.guards; - if (!rules_[id].guard(graph, match)) { - report(id, root->Index(), FailureStage::Guard, - match.RejectionReason().empty() ? std::string_view("Guard rejected match.") : match.RejectionReason()); - continue; + 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; + auto consider = [&](const Pattern& pattern, gsl::span rule_ids) { + statistics.candidates += rule_ids.size(); + ++statistics.structural_checks; + Match match; + FailureStage stage = FailureStage::Structure; + std::string_view reason; + if (!MatchPattern(pattern, graph, *root, producers, match, stage, reason)) { + for (size_t id : rule_ids) report(id, root->Index(), stage, reason); + return; } - if (options_.match_only) ++statistics.matches; - if (id < winner) { - winner = id; - winning_match = match; - winning_pattern = match.pattern; + if (!CanRewritePartition(match, options_)) { + for (size_t id : rule_ids) + report(id, root->Index(), FailureStage::ExecutionProvider, "Mixed or unsupported execution providers."); + return; } - if (!options_.match_only) break; - } - }; - if (options_.indexed) { - auto it = roots_.find(Key(root->OpType(), root->Domain())); - if (it != roots_.end()) { - for (size_t id : it->second) { - const auto& group = groups_[id]; - const auto& ids = applicable_groups[id]; - if (ids.empty() || (!options_.match_only && ids.front() >= winner)) continue; - consider(group.pattern, ids); + for (size_t id : rule_ids) { + if (!options_.match_only && id >= winner) break; + match.pattern = &rules_[id].pattern; + match.ClearUpdates(); + ++statistics.guards; + if (!rules_[id].guard(graph, match)) { + report(id, root->Index(), FailureStage::Guard, + match.RejectionReason().empty() ? std::string_view("Guard rejected match.") : match.RejectionReason()); + continue; + } + if (options_.match_only) ++statistics.matches; + if (id < winner) { + winner = id; + winning_match = match; + winning_pattern = match.pattern; + } + if (!options_.match_only) break; + } + }; + if (options_.indexed) { + auto it = roots_.find(Key(root->OpType(), root->Domain())); + if (it != roots_.end()) { + for (size_t id : it->second) { + const auto& group = groups_[id]; + const auto& ids = applicable_groups[id]; + if (ids.empty() || (!options_.match_only && ids.front() >= winner)) continue; + consider(group.pattern, ids); + } + } + } else { + for (size_t id : applicable_rules) { + consider(rules_[id].pattern, gsl::span(&id, 1)); + if (!options_.match_only && winner != std::numeric_limits::max()) break; } } - } else { - for (size_t id : applicable_rules) { - consider(rules_[id].pattern, gsl::span(&id, 1)); - if (!options_.match_only && winner != std::numeric_limits::max()) 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); } - } - 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; - winning_match.ClearUpdates(); - 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) worklist.push_back(index); - for (auto it = node->OutputEdgesBegin(); it != node->OutputEdgesEnd(); ++it) { - next.push_back(it->GetNode().Index()); + InlinedVector added; + winning_match.ClearUpdates(); + 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()); + } + 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; } } affected = std::move(next); diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 692a6bd9994c1..a2cc86501de70 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -171,6 +171,8 @@ struct Statistics { size_t guards = 0; size_t matches = 0; size_t rewrites = 0; + size_t iterations = 0; + bool iteration_limit_reached = false; }; enum class FailureStage { Opset, @@ -196,6 +198,7 @@ struct Options { 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. diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index bc3585aadf074..af2f5b85e32e7 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -517,7 +517,7 @@ TEST(PatternRewriterTests, ReplacementParsingCompilationAndRollback) { invalid.replacement = Replacement("(a)=>Identity(a)"); EXPECT_THROW(PatternRewriter({invalid}), OnnxRuntimeException); invalid.replacement = Replacement("(x)=>a=Identity(x);return Identity(a)"); - EXPECT_THROW(PatternRewriter({invalid}), OnnxRuntimeException); + 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)")); @@ -1352,11 +1352,118 @@ TEST(PatternRewriterTests, PriorityOverlapTerminationAndMatchOnly) { } } +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.guard = [](const Graph& graph, const 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.guard = [](const Graph&, const 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_THROW(PatternRewriter({rule}), OnnxRuntimeException); + 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(); From f5b6d5ef76c11573c6a6870c3690f8de5d1c0e2f Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 14 Oct 2026 10:38:09 -0700 Subject: [PATCH 14/45] Finish candidate batches outside the descendant invalidation loop Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/optimizer/pattern_rewriter.cc | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 0402308c0bba6..4b7b558b19d66 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -1222,14 +1222,14 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { for (auto it = node->OutputEdgesBegin(); it != node->OutputEdgesEnd(); ++it) { next.push_back(it->GetNode().Index()); } - 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; } + affected = std::move(next); } - 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(); } From 90a830e173faf12d390c810028627b69dcf08e46 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 14 Oct 2026 10:41:39 -0700 Subject: [PATCH 15/45] Keep replacement rollback regression invalid through schema validation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/optimizer/pattern_rewriter_test.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index af2f5b85e32e7..4ca91829f5709 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -648,7 +648,7 @@ TEST(PatternRewriterTests, ReplacementPlanBindingAndSequenceInvalidation) { } auto invalid = second; invalid.replacement = [](const Graph&, const Match&) { - return Replacement("(second)=>a=Identity(second);b=Identity(a);c=Identity(b);return Identity(c)"); + 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); @@ -662,6 +662,7 @@ TEST(PatternRewriterTests, ReplacementPlanBindingAndSequenceInvalidation) { 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()); } From ec4e4730bc517e4ef3fd41e6bc0997bf30941971 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 14 Oct 2026 11:42:05 -0700 Subject: [PATCH 16/45] Outline backward matching and indexed bounded rule scheduling Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 41 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 066be14d3d74f..79380e7484e71 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -525,3 +525,44 @@ Tests in `onnxruntime/test/optimizer/pattern_rewriter_test.cc` cover parsing, guards, 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 +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 semantic guard. + +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 guards + 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 guards; +patterns with attribute predicates are kept separate. The earliest passing +rule wins even across groups, and lower-priority work is skipped when possible. +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. From 7893d8c092f550926e4cd2993a5890edc86b0634 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Thu, 15 Oct 2026 07:24:06 -0700 Subject: [PATCH 17/45] Add tensor type and shape annotations with implicit dimension bindings Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 85 ++++++- .../core/optimizer/pattern_rewriter.cc | 221 ++++++++++++++++- onnxruntime/core/optimizer/pattern_rewriter.h | 26 ++ .../test/optimizer/pattern_rewriter_test.cc | 229 +++++++++++++++++- 4 files changed, 549 insertions(+), 12 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 79380e7484e71..515b14c7047e4 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -103,6 +103,84 @@ values. Unknown names and unbound values raise explicit errors. Read-only `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. +`?` still describes an absent operator slot, not an ONNX optional value type: +`(x:float[], low?:float, high?:float)` allows absent scalar Clip bounds. +An absent slot satisfies its annotation without binding any dimensions or 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. Repeated +known bindings must agree exactly: integer 2 differs from symbolic `"batch"`, +and `"batch"` differs from `"other"`. An anonymous unknown dimension leaves +the variable unbound and does not conflict with a later known binding. +Consequently, repeated names do **not** prove runtime equality when occurrences +are unknown; guards requiring that proof must reject unknown metadata. +Likewise, a missing element type leaves `T` unbound, while a concrete annotation +such as `float` requires a known matching element type. + +In a guard 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, carry +shapes, or be marked optional. Undeclared captures remain supported and untyped. + +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. + ### Optional input slots Use `None` in an operator input position to require an absent input. It matches @@ -184,7 +262,8 @@ Replacement replacement("(x,axes) => Transpose(x,perm=axes)"); The attribute must be present and concrete. `match.attributeOf("axes")` exposes the captured `AttributeProto` to guards and replacements without a callback. -Capture names are separate from tensor values and cannot collide with them. +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 @@ -494,7 +573,7 @@ options.on_failure = [](const MatchFailure& failure) { }; ``` -Stages distinguish opset filtering, structure, constants, attributes, deletion +Stages distinguish opset filtering, structure, type/shape constraints, constants, attributes, deletion boundaries, execution providers, guards, and replacement validation. An opset rejection is reported once per rule/graph with root index `std::numeric_limits::max()` because no candidate was inspected. @@ -561,7 +640,7 @@ 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 guards; -patterns with attribute predicates are kept separate. The earliest passing +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. With `indexed=false`, the reference mode instead tries all opset-applicable rules in registration order at each node. Match-only mode processes one batch, diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 4b7b558b19d66..498f982c81125 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -59,6 +59,46 @@ std::string Signature(const Pattern& pattern) { using Producers = InlinedHashMap; +bool MatchTypes(const Pattern& pattern, Match& match) { + for (const auto& [handle, constraint] : pattern.TypeConstraints()) { + if (!constraint.element_variable.empty() && !match.BindElementType(constraint.element_variable, 0)) return false; + for (const auto& dim : constraint.dimensions) + if (const auto* name = std::get_if(&dim)) + if (!match.BindDimension(*name, std::monostate{})) return false; + if (pattern.IsAttribute(handle)) { + if (match.attributeOf(constraint.value_name).type() != constraint.attribute_type) return false; + continue; + } + const NodeArg* value = match.values[handle]; + if (!value) continue; // An absent optional slot has no tensor metadata to constrain. + 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) continue; + 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))) return false; + } + } + } + return true; +} + bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, const Producers& producers, Match& match, FailureStage& stage, std::string_view& reason) { auto fail = [&](FailureStage failure_stage, std::string_view message) { @@ -131,6 +171,8 @@ bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, } } // Input variables must be outside the deletion set, even when aliasing another binding. + if (!MatchTypes(pattern, match)) + return fail(FailureStage::TypeShape, "Tensor type/rank/dimension or captured attribute type mismatch."); for (int value = 0; value < static_cast(pattern.ValueCount()); ++value) { if (!pattern.IsInput(value)) continue; auto it = producers.find(match.values[value]); @@ -238,7 +280,8 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R value = pattern.NamedValue(std::get(input)); } ORT_RETURN_IF( - value.handle < 0 || static_cast(value.handle) >= match.values.size() || !pattern.IsInput(value), + value.handle < 0 || static_cast(value.handle) >= match.values.size() || !pattern.IsInput(value) || + pattern.IsAttribute(value), "Replacement inputs must be pattern input variables."); const NodeArg* binding = match.values[value]; inputs.push_back(binding ? graph.GetNodeArg(binding->Name()) : &graph.GetOrCreateNodeArg("", nullptr)); @@ -563,6 +606,11 @@ class PatternParser { pattern_.optional_inputs_.insert(input.handle); ORT_RETURN_IF_ERROR(Next()); } + if (token_ == ":") { + TypeConstraint constraint; + ORT_RETURN_IF_ERROR(Annotation(name, constraint)); + pattern_.type_constraints_.emplace(input.handle, std::move(constraint)); + } if (token_ != ",") break; ORT_RETURN_IF_ERROR(Next()); } while (true); @@ -578,6 +626,7 @@ class PatternParser { continue; } InlinedVector names; + InlinedHashMap annotations; if (kind_ == Kind::Identifier && token_ == "return") { returned = true; ORT_RETURN_IF_ERROR(Next()); @@ -590,6 +639,12 @@ class PatternParser { ORT_RETURN_IF_ERROR(Identifier(name, names.empty())); 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")); + TypeConstraint constraint; + ORT_RETURN_IF_ERROR(Annotation(names.back(), constraint)); + annotations.emplace(names.back(), std::move(constraint)); + } if (token_ != ",") break; ORT_RETURN_IF_ERROR(Next()); } while (true); @@ -621,7 +676,11 @@ class PatternParser { } } } - if (!returned_values) ORT_RETURN_IF_ERROR(Call(names)); + if (!returned_values) { + ORT_RETURN_IF_ERROR(Call(names)); + for (auto& [name, constraint] : annotations) + pattern_.type_constraints_.emplace(symbols_.at(name).handle, std::move(constraint)); + } if (returned) { if (!returned_values) pattern_.outputs_ = pattern_.Nodes().back().outputs; if (kind_ == Kind::Symbol && token_ == ";") ORT_RETURN_IF_ERROR(Next()); @@ -631,6 +690,29 @@ class PatternParser { } } ORT_RETURN_IF(!returned, Error("missing return")); + InlinedHashSet dimension_names; + InlinedHashSet element_names; + for (const auto& [handle, constraint] : pattern_.type_constraints_) { + if (pattern_.IsAttribute(handle)) { + ORT_RETURN_IF(!constraint.attribute_type || pattern_.IsOptional(handle), + Error("invalid attribute annotation or optional attribute")); + } else { + ORT_RETURN_IF(!constraint.tensor_type, Error("annotation requires an attribute variable")); + } + if (!constraint.element_variable.empty()) element_names.insert(constraint.element_variable); + for (const auto& dim : constraint.dimensions) + if (const auto* name = std::get_if(&dim)) dimension_names.insert(*name); + } + ORT_RETURN_IF(dimension_names.size() + element_names.size() > 256, Error("too many annotation variables")); + for (const auto& name : dimension_names) + ORT_RETURN_IF(symbols_.count(name) || pattern_.attribute_names_.count(name) || element_names.count(name), + Error("dimension name conflicts with another variable")); + for (const auto& name : element_names) + ORT_RETURN_IF(symbols_.count(name) || pattern_.attribute_names_.count(name), + Error("element type name conflicts with another variable")); + 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; for (Value output : pattern_.Nodes().back().outputs) reachable.insert(output.handle); for (int handle : attribute_inputs_) reachable.insert(handle); @@ -663,7 +745,9 @@ class PatternParser { if (position_ == text_.size()) return Status::OK(); kind_ = Kind::Symbol; char c = text_[position_++]; - if (Alpha(c)) { + 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_] == '.')) @@ -697,7 +781,7 @@ class PatternParser { (c == ':' && position_ < text_.size() && text_[position_] == ':')) { ++position_; } else { - ORT_RETURN_IF(std::string_view("(),=;{}[]?").find(c) == std::string_view::npos, Error("unexpected character")); + 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(); @@ -716,6 +800,80 @@ class PatternParser { return Next(); } + Status Annotation(const std::string& name, TypeConstraint& constraint) { + ORT_RETURN_IF(replacement_attributes_, Error("type annotations are supported in patterns only")); + const size_t start = token_start_; + constraint.value_name = name; + ORT_RETURN_IF_ERROR(Expect(":")); + std::string type; + 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("]")); + } + constraint.spelling = std::string(text_.substr(start, token_start_ - start)); + return Status::OK(); + } + Status Number(std::variant& result) { ORT_RETURN_IF(kind_ != Kind::Number, Error("expected numeric literal")); std::string_view number = token_; @@ -808,7 +966,13 @@ class PatternParser { bindings.emplace(name, capture); attribute_inputs_.insert(found->second.handle); } else { - ORT_RETURN_IF(symbols_.count(capture), Error("attribute capture conflicts with a value name")); + 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); attributes.push_back({name, [](const auto&) { return true; }, capture}); } @@ -894,6 +1058,7 @@ class PatternParser { 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])); } @@ -978,6 +1143,8 @@ 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(std::none_of(pattern->values_.begin(), pattern->values_.end(), [&](const Value& binding) { return binding.name == name; }), @@ -1006,6 +1173,39 @@ bool Match::BindAttribute(std::string_view name, const ONNX_NAMESPACE::Attribute return inserted || it->second.SerializeAsString() == normalized.SerializeAsString(); } +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"), rules_(std::move(rules)), options_(options) { ORT_ENFORCE(options_.max_iterations > 0, "PatternRewriter max_iterations must be positive."); @@ -1031,13 +1231,16 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) ORT_ENFORCE(domains.count(node.domain), "Missing rule opset requirement for replacement domain: ", node.domain); for (const auto& [name, binding] : node.attribute_bindings) { ORT_ENFORCE(std::none_of(rule.pattern.values_.begin(), rule.pattern.values_.end(), - [&](const Value& value) { return value.name == binding; }), + [&](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), "Replacement inputs must be pattern input variables: ", *name); + ORT_ENFORCE(rule.pattern.IsInput(value) && !rule.pattern.IsAttribute(value), + "Replacement inputs must be pattern tensor input variables: ", *name); input = std::move(value); } } @@ -1056,13 +1259,15 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) for (Value input : it->inputs) if (input.handle >= 0) reachable.insert(input.handle); } + for (int handle : rule.pattern.attribute_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(); })) + [](const auto& node) { return !node.attributes.empty(); }) || + !rule.pattern.TypeConstraints().empty()) signature += "#" + std::to_string(i); auto [it, inserted] = signatures.emplace(signature, groups_.size()); if (inserted) { diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index a2cc86501de70..797fa09e01588 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -49,6 +50,20 @@ struct Operator { using Literal = std::variant, InlinedVector>; +using Dimension = std::variant; + +struct TypeConstraint { + std::string value_name; + std::string spelling; + 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; +}; + struct OpsetRequirement { std::string domain; int min_version; @@ -74,6 +89,8 @@ class Pattern { 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; } private: friend class detail::PatternParser; @@ -96,6 +113,8 @@ class Pattern { InlinedVector outputs_; InlinedHashSet optional_inputs_; InlinedHashSet attribute_names_; + InlinedHashMap type_constraints_; + InlinedHashSet attribute_inputs_; Value AddValue(std::string_view name); }; @@ -105,11 +124,15 @@ struct Match { const NodeArg* valueOf(std::string_view name) const; void update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) const; const ONNX_NAMESPACE::AttributeProto& attributeOf(std::string_view name) const; + const Dimension& dimensionOf(std::string_view name) const; + std::optional elementTypeOf(std::string_view name) const; void ClearUpdates() const { extra_attributes_.clear(); rejection_reason_.clear(); } bool BindAttribute(std::string_view name, const ONNX_NAMESPACE::AttributeProto& value); + bool BindDimension(std::string_view name, Dimension value); + bool BindElementType(std::string_view name, int32_t value); bool Reject(std::string_view reason) const { rejection_reason_ = reason; return false; @@ -119,6 +142,8 @@ struct Match { private: NodeAttributes captured_attributes_; + InlinedHashMap dimensions_; + InlinedHashMap element_types_; mutable NodeAttributes extra_attributes_; mutable std::string rejection_reason_; }; @@ -177,6 +202,7 @@ struct Statistics { enum class FailureStage { Opset, Structure, + TypeShape, Constant, Attribute, Boundary, diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 4ca91829f5709..bf1d1b7d2681c 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -701,7 +701,7 @@ TEST(PatternRewriterTests, ReplacementExtraAttributeBindings) { "(x)=>a=Identity(x);return HardSigmoid(a,alpha=a)", "(x,a)=>HardSigmoid(x,alpha=a,alpha=0.5)"}) EXPECT_THROW(Replacement{text}, OnnxRuntimeException); - EXPECT_THROW(Pattern{"(x,a)=>return HardSigmoid(x,alpha=a)"}, 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"); @@ -2164,6 +2164,233 @@ TEST(PatternRewriterTests, StructuredFailureDiagnostics) { } } +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[...])=>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.guard = [scenario, &guards](const Graph&, const 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"); + if (scenario == 5) EXPECT_TRUE(std::holds_alternative(match.dimensionOf("M"))); + 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 || scenario == 5; + 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,hi?:float)=>{a=Clip(x,lo,hi);return Identity(a);}"); + second.pattern = Pattern( + "(input:tensor[M,N],low?:tensor,high?:float)=>" + "{a:tensor[M,N]=Clip(input,low,high);return Identity(a);}"); + second.replacement = Replacement("(input)=>Clip(input)"); + second.guard = [](const Graph&, const 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,hi?:float)=>" + "{a=Clip(x,lo,hi);return Identity(a);}"); + if (scenario == 1) first.guard = [](const Graph&, const Match&) { return false; }; + if (scenario == 2) second.pattern = Pattern( + "(input:tensor[M,M],low?:tensor,high?:float)=>" + "{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, 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"); From 328ef440a83a33a14c1408c35ce489cf711d9d4f Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Thu, 15 Oct 2026 08:14:09 -0700 Subject: [PATCH 18/45] Support bounded type unions and express optional inputs with None Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 50 ++++-- .../core/optimizer/pattern_rewriter.cc | 166 +++++++++++------- onnxruntime/core/optimizer/pattern_rewriter.h | 4 +- .../test/optimizer/pattern_rewriter_test.cc | 162 +++++++++++++++-- 4 files changed, 300 insertions(+), 82 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 515b14c7047e4..b11a0acd5a223 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -136,9 +136,10 @@ They cannot collide with tensor, attribute, or element-type variable names. 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. -`?` still describes an absent operator slot, not an ONNX optional value type: -`(x:float[], low?:float, high?:float)` allows absent scalar Clip bounds. -An absent slot satisfies its annotation without binding any dimensions or type. +`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 @@ -173,7 +174,7 @@ 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, carry -shapes, or be marked optional. Undeclared captures remain supported and untyped. +shapes, or use type unions (including `None`). Undeclared captures remain supported and untyped. Annotations are constraints on **patterns only**; replacements continue to use unannotated inputs and derive types/shapes from preserved outputs and schema @@ -181,20 +182,46 @@ 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)"); +``` + +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. Guards run once after a consistent combination is selected; +guard rejection does not trigger another combination. + +This bounded search applies only to type/shape annotations, not graph structure: +there is still no search over producers or operand permutations. 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, but `None` is permitted only for header tensor +inputs, not assigned outputs or attribute captures. 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 `?` to allow either a present value or an absent slot: +variable with a tensor type union including `None` to allow either a present value or an absent slot: ```cpp -Pattern pattern("(x,low?,high?) => { a=Identity(x); return Clip(a,low,high); }"); -Replacement replacement("(x,low?,high?) => Clip(x,low,high)"); +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?) => Clip(x,v,v)` does not match +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; guards must check it before dereferencing. Replacements preserve an absent binding as an empty input name. Explicit `None` @@ -595,7 +622,7 @@ normally and are not converted to nonmatches. Definitions are limited to 64 nodes, 256 values, and 64 inputs/outputs per operator. They are connected acyclic graphs rooted at the final operator. -There is no alternation, implicit commutation, variadic matching, +There is no 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`. @@ -608,7 +635,7 @@ 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 -backtracking. Starting with the pattern's final operator bound to that model +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 @@ -616,7 +643,8 @@ 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 semantic guard. +before the rule's semantic guard. Type unions use the separate bounded +annotation-choice search described above. Multiple rules are handled **node-first**, inside an outer bounded retry loop: diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 498f982c81125..75236033032fd 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -59,44 +59,64 @@ std::string Signature(const Pattern& pattern) { using Producers = InlinedHashMap; +bool MatchType(const Pattern& pattern, int handle, const TypeConstraint& constraint, Match& match) { + if (pattern.IsAttribute(handle)) { + return match.attributeOf(constraint.value_name).type() == constraint.attribute_type; + } + const NodeArg* value = match.values[handle]; + if (constraint.absent) return !value; + if (!value) 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))) return false; + } + } + return true; +} + bool MatchTypes(const Pattern& pattern, Match& match) { - for (const auto& [handle, constraint] : pattern.TypeConstraints()) { - if (!constraint.element_variable.empty() && !match.BindElementType(constraint.element_variable, 0)) return false; - for (const auto& dim : constraint.dimensions) - if (const auto* name = std::get_if(&dim)) - if (!match.BindDimension(*name, std::monostate{})) return false; - if (pattern.IsAttribute(handle)) { - if (match.attributeOf(constraint.value_name).type() != constraint.attribute_type) return false; - continue; + InlinedVector handles; + for (const auto& [handle, alternatives] : pattern.TypeConstraints()) { + handles.push_back(handle); + 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{}); } - const NodeArg* value = match.values[handle]; - if (!value) continue; // An absent optional slot has no tensor metadata to constrain. - 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) continue; - 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))) return false; + } + std::sort(handles.begin(), handles.end()); + std::function search = [&](size_t index, Match& state) { + if (index == handles.size()) return true; + const int handle = handles[index]; + for (const auto& constraint : pattern.TypeConstraints().at(handle)) { + Match candidate = state; + if (MatchType(pattern, handle, constraint, candidate) && search(index + 1, candidate)) { + state = std::move(candidate); + return true; } } - } - return true; + return false; + }; + return search(0, match); } bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, const Producers& producers, @@ -602,14 +622,12 @@ class PatternParser { ORT_RETURN_IF(pattern_.ValueCount() >= 256, Error("too many values")); Value input = pattern_.Input(name); symbols_.emplace(name, input); - if (token_ == "?") { - pattern_.optional_inputs_.insert(input.handle); - ORT_RETURN_IF_ERROR(Next()); - } if (token_ == ":") { - TypeConstraint constraint; - ORT_RETURN_IF_ERROR(Annotation(name, constraint)); - pattern_.type_constraints_.emplace(input.handle, std::move(constraint)); + 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_.type_constraints_.emplace(input.handle, std::move(alternatives)); } if (token_ != ",") break; ORT_RETURN_IF_ERROR(Next()); @@ -626,7 +644,7 @@ class PatternParser { continue; } InlinedVector names; - InlinedHashMap annotations; + InlinedHashMap> annotations; if (kind_ == Kind::Identifier && token_ == "return") { returned = true; ORT_RETURN_IF_ERROR(Next()); @@ -641,9 +659,11 @@ class PatternParser { names.push_back(std::move(name)); if (token_ == ":") { ORT_RETURN_IF(returned, Error("annotate the definition, not the return")); - TypeConstraint constraint; - ORT_RETURN_IF_ERROR(Annotation(names.back(), constraint)); - annotations.emplace(names.back(), std::move(constraint)); + 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()); @@ -692,16 +712,21 @@ class PatternParser { ORT_RETURN_IF(!returned, Error("missing return")); InlinedHashSet dimension_names; InlinedHashSet element_names; - for (const auto& [handle, constraint] : pattern_.type_constraints_) { - if (pattern_.IsAttribute(handle)) { - ORT_RETURN_IF(!constraint.attribute_type || pattern_.IsOptional(handle), - Error("invalid attribute annotation or optional attribute")); - } else { - ORT_RETURN_IF(!constraint.tensor_type, Error("annotation requires an attribute variable")); + size_t combinations = 1; + for (const auto& [handle, alternatives] : pattern_.type_constraints_) { + ORT_RETURN_IF(combinations > 64 / alternatives.size(), Error("type alternatives exceed 64 combinations")); + combinations *= alternatives.size(); + for (const auto& constraint : alternatives) { + if (pattern_.IsAttribute(handle)) { + ORT_RETURN_IF(alternatives.size() != 1 || !constraint.attribute_type || pattern_.IsOptional(handle), + Error("invalid attribute annotation or optional attribute")); + } else { + ORT_RETURN_IF(!constraint.tensor_type, Error("annotation requires an attribute variable")); + } + if (!constraint.element_variable.empty()) element_names.insert(constraint.element_variable); + for (const auto& dim : constraint.dimensions) + if (const auto* name = std::get_if(&dim)) dimension_names.insert(*name); } - if (!constraint.element_variable.empty()) element_names.insert(constraint.element_variable); - for (const auto& dim : constraint.dimensions) - if (const auto* name = std::get_if(&dim)) dimension_names.insert(*name); } ORT_RETURN_IF(dimension_names.size() + element_names.size() > 256, Error("too many annotation variables")); for (const auto& name : dimension_names) @@ -781,7 +806,7 @@ class PatternParser { (c == ':' && position_ < text_.size() && text_[position_] == ':')) { ++position_; } else { - ORT_RETURN_IF(std::string_view("(),=;{}[]?:<>").find(c) == std::string_view::npos, Error("unexpected character")); + 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(); @@ -800,11 +825,35 @@ class PatternParser { return Next(); } - Status Annotation(const std::string& name, TypeConstraint& constraint) { + Status Annotation(const std::string& name, InlinedVector& alternatives) { ORT_RETURN_IF(replacement_attributes_, Error("type annotations are supported in patterns only")); - const size_t start = token_start_; - constraint.value_name = name; 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; + }; + 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 tensor alternative")); + return Status::OK(); + } + + Status AnnotationType(TypeConstraint& constraint) { std::string type; ORT_RETURN_IF_ERROR(Identifier(type)); if (type == "int" || type == "ints" || type == "floats" || type == "strings") { @@ -870,7 +919,6 @@ class PatternParser { } ORT_RETURN_IF_ERROR(Expect("]")); } - constraint.spelling = std::string(text_.substr(start, token_start_ - start)); return Status::OK(); } diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 797fa09e01588..a015b0c6e19f4 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -54,7 +54,6 @@ using Dimension = std::variant; struct TypeConstraint { std::string value_name; - std::string spelling; int32_t element_type = 0; // Zero means unconstrained. std::string element_variable; InlinedVector dimensions; @@ -62,6 +61,7 @@ struct TypeConstraint { bool prefix = false; int32_t attribute_type = 0; // Scalar/list attribute interpretation, when applicable. bool tensor_type = true; + bool absent = false; }; struct OpsetRequirement { @@ -113,7 +113,7 @@ class Pattern { InlinedVector outputs_; InlinedHashSet optional_inputs_; InlinedHashSet attribute_names_; - InlinedHashMap type_constraints_; + InlinedHashMap> type_constraints_; InlinedHashSet attribute_inputs_; Value AddValue(std::string_view name); }; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index bf1d1b7d2681c..11287602741ef 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1973,11 +1973,14 @@ TEST(PatternRewriterTests, NamedAttributeCaptures) { } TEST(PatternRewriterTests, OptionalInputSlotsAndNumerics) { - for (auto text : {"(x,low?,high?)=>{a=Identity(x);return Clip(a,low,high);}", + 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}); - EXPECT_NO_THROW((void)Replacement{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);}", @@ -1987,8 +1990,8 @@ TEST(PatternRewriterTests, OptionalInputSlotsAndNumerics) { for (int scenario = 0; scenario < 5; ++scenario) { SCOPED_TRACE(scenario); auto rule = IdentityPair(); - rule.pattern = Pattern("(x,low?,high?)=>{a=Identity(x);return Clip(a,low,high);}"); - rule.replacement = Replacement("(x,low?,high?)=>Clip(x,low,high)"); + 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(); @@ -2103,7 +2106,7 @@ TEST(PatternRewriterTests, OptionalInputConstraintsAndAbsentRepetition) { ASSERT_STATUS_OK(test.graph.Resolve()); for (int constraint = 0; constraint < 3; ++constraint) { auto rule = IdentityPair(); - rule.pattern = Pattern(constraint == 0 ? "(x,v?)=>{a=Identity(x);return Clip(a,v,v);}" + 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; @@ -2172,7 +2175,7 @@ TEST(PatternRewriterTests, TypeShapeAnnotationParsing) { "(x:float)=>return Identity(x)", "(x:tensor[M,0])=>return Identity(x)", "(x:tensor[M,...],y:tensor[M,_])=>return CastLike(x,y)", - "(x?:int64[...])=>return Identity(x)", + "(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)", @@ -2359,9 +2362,9 @@ TEST(PatternRewriterTests, AnnotatedOptionalValuesAndRuleIsolation) { b.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); auto first = IdentityPair("reject"), second = IdentityPair("accept"); - first.pattern = Pattern("(x:float[2,4],lo?:float,hi?:float)=>{a=Clip(x,lo,hi);return Identity(a);}"); + 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,high?:float)=>" + "(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.guard = [](const Graph&, const Match& match) { @@ -2371,11 +2374,11 @@ TEST(PatternRewriterTests, AnnotatedOptionalValuesAndRuleIsolation) { return true; }; if (scenario == 1) first.pattern = Pattern( - "(x:float[2,3],lo?:float,hi?:float)=>" + "(x:float[2,3],lo:float | None,hi:float | None)=>" "{a=Clip(x,lo,hi);return Identity(a);}"); if (scenario == 1) first.guard = [](const Graph&, const Match&) { return false; }; if (scenario == 2) second.pattern = Pattern( - "(input:tensor[M,M],low?:tensor,high?:float)=>" + "(input:tensor[M,M],low:tensor | None,high:float | None)=>" "{a=Clip(input,low,high);return Identity(a);}"); InlinedVector failures; Options options{.indexed = indexed}; @@ -2391,6 +2394,145 @@ TEST(PatternRewriterTests, AnnotatedOptionalValuesAndRuleIsolation) { } } +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,p:int|None)=>return Identity(x,axis=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.guard = [scenario](const Graph&, const 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.guard = [scenario](const Graph&, const 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.guard = [scenario](const Graph&, const 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.guard = [](const Graph&, const 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"); From eb1fff20381dfea75c27eb9c0b71ea702678b1c3 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Thu, 15 Oct 2026 12:23:53 -0700 Subject: [PATCH 19/45] Simplify rewrite rule construction and reorder parameters Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 27 ++-- .../core/optimizer/pattern_rewriter.cc | 20 +++ onnxruntime/core/optimizer/pattern_rewriter.h | 18 ++- .../core/optimizer/pattern_rewriter_rules.cc | 36 ++--- .../test/optimizer/pattern_rewriter_test.cc | 128 +++++++++++++++--- 5 files changed, 180 insertions(+), 49 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index b11a0acd5a223..267f6188bc5b1 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -6,9 +6,13 @@ operator patterns with semantically equivalent graphs. Include text; programmatic builders and mutable replacement recipes are not public APIs. Guards and replacement callbacks remain ordinary C++ functions. -A `Rule` combines a diagnostic name, a parsed pattern, domain/opset requirements, -a semantic guard, and a parsed replacement or callback returning one. The matcher -establishes structural and deletion safety. **The rule author must establish +A `Rule` takes `(name, pattern, replacement, guard, 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.** ## Example: HardSigmoid followed by Mul @@ -29,18 +33,17 @@ bool IsFloat32(const NodeArg& value) { } Rule MakeHardSwishRule() { - Pattern pattern(R"( + const auto* pattern = R"( (x) => { sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) return Mul(sigmoid, x) } - )"); + )"; auto guard = [](const Graph&, const Match& match) { return IsFloat32(*match.valueOf("x")); }; - Replacement replacement("(x) => HardSwish(x)"); - return Rule{"HardSigmoidMul", std::move(pattern), - {{kOnnxDomain, 14, 21}}, guard, replacement}; + const auto* replacement = "(x) => HardSwish(x)"; + return Rule{"HardSigmoidMul", pattern, replacement, guard, {{kOnnxDomain, 14, 21}}}; } ``` @@ -371,7 +374,8 @@ back; a mapping that preserves an acyclic dependency graph is allowed. Reuse a parsed replacement plan when an attribute depends on each match. Declare the attribute binding in the replacement's input list, set it with -`match.update`, and reference it as a keyword argument. +`match.update`, and reference it as a keyword argument. Parse the plan outside +the callback so matching does not parse it again. The following rule folds division by a finite float32 scalar into FusedMatMul's `alpha` attribute. It requires the Microsoft-domain opset to be imported already. @@ -395,7 +399,7 @@ std::optional FloatScalar(const Graph& graph, const NodeArg& value) { } Rule MakeMatMulDivRule() { - Pattern pattern("(x,y,divisor) => { mm = MatMul(x,y); return Div(mm,divisor); }"); + const auto* pattern = "(x,y,divisor) => { mm = MatMul(x,y); return Div(mm,divisor); }"; auto guard = [](const Graph& graph, const Match& match) { auto divisor = FloatScalar(graph, *match.valueOf("divisor")); return IsFloat32(*match.valueOf("x")) && IsFloat32(*match.valueOf("y")) && @@ -409,8 +413,7 @@ Rule MakeMatMulDivRule() { match.update("alpha", std::move(alpha)); return plan; }; - return Rule{"MatMulDiv", std::move(pattern), - {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}, guard, replacement}; + return Rule{"MatMulDiv", pattern, replacement, guard, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; } ``` diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 75236033032fd..3f10ca72fc6a3 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -1133,6 +1133,26 @@ class PatternParser { }; } // namespace detail +Rule::Rule(std::string name, Pattern pattern, ReplacementFunction replacement, Guard guard, + InlinedVector opsets) + : name(std::move(name)), + pattern(std::move(pattern)), + replacement(std::move(replacement)), + guard(std::move(guard)), + opsets(std::move(opsets)) {} + +Rule::Rule(std::string name, std::string_view pattern, ReplacementFunction replacement, Guard guard, + InlinedVector opsets) + : Rule(std::move(name), Pattern(pattern), std::move(replacement), std::move(guard), std::move(opsets)) {} + +Rule::Rule(std::string name, Pattern pattern, std::string_view replacement, Guard guard, + InlinedVector opsets) + : Rule(std::move(name), std::move(pattern), Replacement(replacement), std::move(guard), std::move(opsets)) {} + +Rule::Rule(std::string name, std::string_view pattern, std::string_view replacement, Guard guard, + InlinedVector opsets) + : Rule(std::move(name), Pattern(pattern), Replacement(replacement), std::move(guard), std::move(opsets)) {} + Pattern::Pattern(std::string_view text) { ORT_THROW_IF_ERROR(detail::PatternParser(text).Parse(*this)); } Replacement::Replacement(std::string_view text) { diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index a015b0c6e19f4..4c3af113ee9fc 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -178,13 +178,25 @@ struct Replacement { }; struct Rule { + using Guard = std::function; + using ReplacementFunction = std::function; + + Rule(std::string name, Pattern pattern, ReplacementFunction replacement, Guard guard, + InlinedVector opsets); + Rule(std::string name, std::string_view pattern, ReplacementFunction replacement, Guard guard, + InlinedVector opsets); + Rule(std::string name, Pattern pattern, std::string_view replacement, Guard guard, + InlinedVector opsets); + Rule(std::string name, std::string_view pattern, std::string_view replacement, Guard guard, + InlinedVector opsets); + std::string name; Pattern pattern; - InlinedVector opsets; // Guards must be pure. Recipes may set extra attribute bindings, but must not mutate the graph // or retain pointers into a Match. - std::function guard; - std::function replacement; + ReplacementFunction replacement; + Guard guard; + InlinedVector opsets; }; using Rules = InlinedVector; diff --git a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc index 4a7049fb0c76c..05be0d38a86cb 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc @@ -47,34 +47,34 @@ bool Integer(const Node& node, const char* name, int64_t value) { Rules MakeExampleRules() { Rules rules; { - Pattern p(R"( + 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 guard = [](const Graph&, const Match& m) { return IsFloat32(*m.valueOf("x")); }; - Replacement replacement("(x) => HardSwish(x)"); - Rule rule{"HardSwish", p, {{kOnnxDomain, 14, 21}}, guard, replacement}; + const auto* replacement = "(x) => HardSwish(x)"; + Rule rule{"HardSwish", pattern, replacement, guard, {{kOnnxDomain, 14, 21}}}; rules.push_back(std::move(rule)); } { - Pattern p(R"( + const auto* pattern = R"( (x) => { sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) return Mul(sigmoid, x) } - )"); + )"; auto guard = [](const Graph&, const Match& m) { return IsFloat32(*m.valueOf("x")); }; - Replacement replacement("(x) => HardSwish(x)"); - Rule rule{"HardSigmoidMul", p, {{kOnnxDomain, 14, 21}}, guard, replacement}; + const auto* replacement = "(x) => HardSwish(x)"; + Rule rule{"HardSigmoidMul", pattern, replacement, guard, {{kOnnxDomain, 14, 21}}}; rules.push_back(std::move(rule)); } #ifndef DISABLE_CONTRIB_OPS { - Pattern p("(x, y, divisor) => { mm = MatMul(x, y); return Div(mm, divisor); }"); + const auto* pattern = "(x, y, divisor) => { mm = MatMul(x, y); return Div(mm, divisor); }"; auto guard = [](const Graph& graph, const Match& m) { auto scalar = Scalar(graph, *m.valueOf("divisor")); return IsFloat32(*m.valueOf("x")) && IsFloat32(*m.valueOf("y")) && @@ -85,11 +85,11 @@ Rules MakeExampleRules() { m.update("alpha", FloatAttribute("alpha", 1.0f / *Scalar(graph, *m.valueOf("divisor")))); return plan; }; - Rule rule{"MatMulDiv", p, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}, guard, replacement}; + Rule rule{"MatMulDiv", pattern, replacement, guard, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; rules.push_back(std::move(rule)); } { - Pattern p("(x, y) => { transpose = Transpose(y); return MatMul(x, transpose); }"); + const auto* pattern = "(x, y) => { transpose = Transpose(y); return MatMul(x, transpose); }"; auto guard = [](const Graph&, const Match& m) { if (!IsFloat32(*m.valueOf("x")) || !IsFloat32(*m.valueOf("y"))) return false; const auto* shape = m.valueOf("y")->Shape(); @@ -107,12 +107,12 @@ Rules MakeExampleRules() { } return true; }; - Replacement replacement("(x,y) => com.microsoft.FusedMatMul(x,y,transB=1)"); - Rule rule{"TransposeMatMul", p, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}, guard, replacement}; + const auto* replacement = "(x,y) => com.microsoft.FusedMatMul(x,y,transB=1)"; + Rule rule{"TransposeMatMul", pattern, replacement, guard, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; rules.push_back(std::move(rule)); } { - Pattern p("(x, y, divisor) => { mm = com.microsoft::FusedMatMul(x, y); return Div(mm, divisor); }"); + const auto* pattern = "(x, y, divisor) => { mm = com.microsoft::FusedMatMul(x, y); return Div(mm, divisor); }"; auto guard = [](const Graph& graph, const Match& m) { auto scalar = Scalar(graph, *m.valueOf("divisor")); auto it = m.nodes[0]->GetAttributes().find("alpha"); @@ -138,11 +138,11 @@ Rules MakeExampleRules() { } return plan; }; - Rule rule{"FusedMatMulDiv", p, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}, guard, replacement}; + Rule rule{"FusedMatMulDiv", pattern, replacement, guard, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; rules.push_back(std::move(rule)); } { - Pattern p(R"( + const auto* pattern = R"( (x, scale, epsilon) => { square = Pow(x, 2.0) mean = ReduceMean(square) @@ -152,7 +152,7 @@ Rules MakeExampleRules() { normalize = Mul(x, reciprocal) return Mul(normalize, scale) } - )"); + )"; auto guard = [](const Graph& graph, const Match& m) { if (!IsFloat32(*m.valueOf("x")) || !IsFloat32(*m.valueOf("scale"))) return false; auto eps = Scalar(graph, *m.valueOf("epsilon")); @@ -179,7 +179,7 @@ Rules MakeExampleRules() { m.update("fused_epsilon", FloatAttribute("epsilon", *Scalar(graph, *m.valueOf("epsilon")))); return plan; }; - Rule rule{"RmsNormalization", p, {{kOnnxDomain, 7, 17}}, guard, replacement}; + Rule rule{"RmsNormalization", pattern, replacement, guard, {{kOnnxDomain, 7, 17}}}; rules.push_back(std::move(rule)); } #endif diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 11287602741ef..48f96bdf2de6e 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -75,15 +75,15 @@ size_t Count(const Graph& graph, const std::string& op) { } Rule IdentityPair(const char* name = "identity", bool accept = true) { - Pattern p("(x) => a = Identity(x); return Identity(a)"); + const auto* pattern = "(x) => a = Identity(x); return Identity(a)"; auto guard = [accept](const Graph&, const Match&) { return accept; }; - Replacement replacement("(x) => Identity(x)"); - Rule rule{name, p, {{kOnnxDomain, 1, 21}}, guard, replacement}; + const auto* replacement = "(x) => Identity(x)"; + Rule rule{name, pattern, replacement, guard, {{kOnnxDomain, 1, 21}}}; return rule; } Rule IdentitySplit() { - Pattern p("(x) => { first = Identity(x); a, b = Split(first); return a, b; }"); + const auto* pattern = "(x) => { first = Identity(x); a, b = Split(first); return a, b; }"; auto guard = [](const Graph&, const Match&) { return true; }; Replacement plan("(x, axis) => { a, b = Split(x, axis=axis); return a, b; }"); auto replacement = [plan](const Graph&, const Match& match) { @@ -93,26 +93,26 @@ Rule IdentitySplit() { : axis->second); return plan; }; - Rule rule{"IdentitySplit", p, {{kOnnxDomain, 13, 17}}, guard, replacement}; + Rule rule{"IdentitySplit", pattern, replacement, guard, {{kOnnxDomain, 13, 17}}}; return rule; } Rule SplitConcat() { - Pattern p("(x) => a, b = Split(x); return Concat(a, b)"); + const auto* pattern = "(x) => a, b = Split(x); return Concat(a, b)"; auto guard = [](const Graph&, const 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; }; - Replacement replacement("(x) => Identity(x)"); - Rule rule{"SplitConcat", p, {{kOnnxDomain, 13, 17}}, guard, replacement}; + const auto* replacement = "(x) => Identity(x)"; + Rule rule{"SplitConcat", pattern, replacement, guard, {{kOnnxDomain, 13, 17}}}; return rule; } // Multi-output adaptation of ONNXScript's LayerNormBiasFusion; unused statistics may be removed. Rule LayerNormBias() { - Pattern p("(x, scale, bias) => y, mean, inv_std = LayerNormalization(x, scale); return Add(y, bias)"); + const auto* pattern = "(x, scale, bias) => y, mean, inv_std = LayerNormalization(x, scale); return Add(y, bias)"; auto guard = [](const Graph&, const Match& match) { const auto* shape = match.valueOf("x")->Shape(); const auto* scale_shape = match.valueOf("scale")->Shape(); @@ -142,7 +142,7 @@ Rule LayerNormBias() { } return plan; }; - Rule rule{"LayerNormBias", p, {{kOnnxDomain, 17, 17}}, guard, replacement}; + Rule rule{"LayerNormBias", pattern, replacement, guard, {{kOnnxDomain, 17, 17}}}; return rule; } @@ -1618,6 +1618,102 @@ TEST(PatternRewriterTests, ExampleRuleImportBoundaries) { } } +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 guard = [](const Graph&, const 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)", guard, {{kOnnxDomain, 1, 21}}}, + {"text", pattern_text, replacement_text, guard, {{kOnnxDomain, 1, 21}}}, + {"parsed", parsed_pattern, parsed_replacement, guard, {{kOnnxDomain, 1, 21}}}, + {"parsed_pattern", parsed_pattern, replacement_text, guard, {{kOnnxDomain, 1, 21}}}, + {"parsed_replacement", pattern_text, parsed_replacement, guard, {{kOnnxDomain, 1, 21}}}, + {"callback", pattern_text, replacement, guard, {{kOnnxDomain, 1, 21}}}, + {"parsed_callback", parsed_pattern, replacement, guard, {{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 guard = [](const Graph&, const Match&) { return true; }; + Rule rule{"owned", pattern, replacement, guard, {{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(); + 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(), plan); + } + size_t calls = 0; + auto callback = [recipe = *compiled, &calls](const Graph&, const Match&) { + ++calls; + return recipe; + }; + Rule computed{"computed", rule.pattern, callback, guard, rule.opsets}; + EXPECT_EQ(calls, 0u); + for (int i = 0; i < 3; ++i) { + auto result = computed.replacement(test.graph, match); + EXPECT_EQ(&result.Plan(), plan); + } + EXPECT_EQ(calls, 3u); +} + +TEST(PatternRewriterTests, RuleConstructorRejectsMalformedText) { + const auto* pattern = "(x)=>{a=Identity(x);return Identity(a);}"; + const auto* replacement = "(x)=>Identity(x)"; + auto guard = [](const Graph&, const Match&) { return true; }; + auto make_rule = [&](auto pattern, auto replacement) { + return Rule{"bad", pattern, replacement, guard, {{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; @@ -1635,10 +1731,10 @@ TEST(PatternRewriterTests, PriorityAcrossInterleavedStructuralGroups) { lower_priority.guard = [](const Graph& graph, const Match& m) { return graph.NodeProducesGraphOutput(*m.nodes.back()); }; - Pattern p("(x)=>a=Identity(x);b=Identity(a);return Identity(b)"); + const auto* pattern = "(x)=>a=Identity(x);b=Identity(a);return Identity(b)"; auto guard = [](const Graph&, const Match&) { return true; }; - Replacement replacement("(x)=>Identity(x)"); - Rule higher_priority{"higher_priority", p, {{kOnnxDomain, 1, 21}}, guard, replacement}; + const auto* replacement = "(x)=>Identity(x)"; + Rule higher_priority{"higher_priority", pattern, replacement, guard, {{kOnnxDomain, 1, 21}}}; Statistics stats; ASSERT_STATUS_OK( PatternRewriter({rejected, higher_priority, lower_priority}, {.indexed = indexed}).Run(test.graph, stats)); @@ -2912,10 +3008,10 @@ TEST(PatternRewriterTests, BranchedSharedNodeBindings) { if (external) b.AddNode("Identity", {sum}, {b.MakeOutput()}); b.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); - Pattern p("(x,y)=>sum=Add(x,y);product=Mul(sum,x);return Add(sum,product)"); + const auto* pattern = "(x,y)=>sum=Add(x,y);product=Mul(sum,x);return Add(sum,product)"; auto guard = [](const Graph&, const Match&) { return true; }; - Replacement replacement("(x)=>Identity(x)"); - Rule rule{"diamond", p, {{kOnnxDomain, 7, 21}}, guard, replacement}; + const auto* replacement = "(x)=>Identity(x)"; + Rule rule{"diamond", pattern, replacement, guard, {{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)); From b330a56cf4e615772c3bcded19bc242a7fb70d29 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Thu, 15 Oct 2026 12:27:01 -0700 Subject: [PATCH 20/45] Compare shared replacement recipes rather than pointer wrappers Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxruntime/test/optimizer/pattern_rewriter_test.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 48f96bdf2de6e..d0ec2cbd5871d 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1674,7 +1674,7 @@ TEST(PatternRewriterTests, RuleConstructorOwnsParsedTextAndReusesPlans) { EXPECT_EQ(rule.pattern.Nodes().size(), 2u); const auto* compiled = rule.replacement.target(); ASSERT_NE(compiled, nullptr); - const auto* plan = &compiled->Plan(); + const auto* plan = compiled->Plan().get(); auto copied = rule; auto moved = std::move(copied); copied = moved; @@ -1683,7 +1683,7 @@ TEST(PatternRewriterTests, RuleConstructorOwnsParsedTextAndReusesPlans) { Match match; for (int i = 0; i < 3; ++i) { auto result = copied.replacement(test.graph, match); - EXPECT_EQ(&result.Plan(), plan); + EXPECT_EQ(result.Plan().get(), plan); } size_t calls = 0; auto callback = [recipe = *compiled, &calls](const Graph&, const Match&) { @@ -1694,7 +1694,7 @@ TEST(PatternRewriterTests, RuleConstructorOwnsParsedTextAndReusesPlans) { EXPECT_EQ(calls, 0u); for (int i = 0; i < 3; ++i) { auto result = computed.replacement(test.graph, match); - EXPECT_EQ(&result.Plan(), plan); + EXPECT_EQ(result.Plan().get(), plan); } EXPECT_EQ(calls, 3u); } From 09ac206669bcb3ed3b48b842b1d634ebbe18ec30 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Thu, 15 Oct 2026 14:26:02 -0700 Subject: [PATCH 21/45] Combine rewrite checks and match-local binding preparation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 109 ++++--- .../core/optimizer/pattern_rewriter.cc | 40 +-- onnxruntime/core/optimizer/pattern_rewriter.h | 30 +- .../core/optimizer/pattern_rewriter_rules.cc | 77 +++-- .../test/optimizer/pattern_rewriter_test.cc | 300 ++++++++++++------ 5 files changed, 328 insertions(+), 228 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 267f6188bc5b1..e3638adf80a59 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -4,9 +4,9 @@ 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. -Guards and replacement callbacks remain ordinary C++ functions. +Check-and-bind callbacks and replacement callbacks remain ordinary C++ functions. -A `Rule` takes `(name, pattern, replacement, guard, opset_requirements)`. +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; @@ -39,16 +39,16 @@ Rule MakeHardSwishRule() { return Mul(sigmoid, x) } )"; - auto guard = [](const Graph&, const Match& match) { + auto check_and_bind = [](const Graph&, Match& match) { return IsFloat32(*match.valueOf("x")); }; const auto* replacement = "(x) => HardSwish(x)"; - return Rule{"HardSigmoidMul", pattern, replacement, guard, {{kOnnxDomain, 14, 21}}}; + return Rule{"HardSigmoidMul", pattern, replacement, check_and_bind, {{kOnnxDomain, 14, 21}}}; } ``` Both occurrences of `x` must bind the **same `NodeArg`**, not just tensors with -equal shapes or contents. The guard checks the relevant data binding `x`; +equal shapes or contents. The check-and-bind callback checks the relevant data binding `x`; resolved schemas establish the corresponding intermediate and output types. It does not require every matched value to have the same element type. @@ -152,11 +152,11 @@ known bindings must agree exactly: integer 2 differs from symbolic `"batch"`, and `"batch"` differs from `"other"`. An anonymous unknown dimension leaves the variable unbound and does not conflict with a later known binding. Consequently, repeated names do **not** prove runtime equality when occurrences -are unknown; guards requiring that proof must reject unknown metadata. +are unknown; check-and-bind callbacks requiring that proof must reject unknown metadata. Likewise, a missing element type leaves `T` unbound, while a concrete annotation such as `float` requires a known matching element type. -In a guard or replacement callback, `match.dimensionOf("M")` returns a +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 @@ -199,8 +199,8 @@ 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. Guards run once after a consistent combination is selected; -guard rejection does not trigger another combination. +remain unbound. Check-and-bind callbacks run once after a consistent combination is selected; +Callback rejection does not trigger another combination. This bounded search applies only to type/shape annotations, not graph structure: there is still no search over producers or operand permutations. The product @@ -226,7 +226,7 @@ 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; guards must check it before dereferencing. +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 @@ -280,7 +280,7 @@ 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++ guard for ranges, shapes, or other semantic conditions. +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: @@ -291,7 +291,7 @@ Replacement replacement("(x,axes) => Transpose(x,perm=axes)"); ``` The attribute must be present and concrete. `match.attributeOf("axes")` exposes -the captured `AttributeProto` to guards and replacements without a callback. +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 @@ -370,12 +370,14 @@ 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. -## Computed attributes with a replacement callback +## Checking applicability and preparing bindings -Reuse a parsed replacement plan when an attribute depends on each match. +`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. Parse the plan outside -the callback so matching does not parse it again. +`match.update`, and reference it as a keyword argument in a text replacement. +No replacement lambda is needed just to prepare attributes. The following rule folds division by a finite float32 scalar into FusedMatMul's `alpha` attribute. It requires the Microsoft-domain opset to be imported already. @@ -400,35 +402,44 @@ std::optional FloatScalar(const Graph& graph, const NodeArg& value) { Rule MakeMatMulDivRule() { const auto* pattern = "(x,y,divisor) => { mm = MatMul(x,y); return Div(mm,divisor); }"; - auto guard = [](const Graph& graph, const Match& match) { + auto check_and_bind = [](const Graph& graph, Match& match) { auto divisor = FloatScalar(graph, *match.valueOf("divisor")); - return IsFloat32(*match.valueOf("x")) && IsFloat32(*match.valueOf("y")) && - divisor && *divisor != 0.0f && std::isfinite(1.0f / *divisor); - }; - Replacement plan("(x,y,alpha) => com.microsoft::FusedMatMul(x,y,alpha=alpha)"); - auto replacement = [plan](const Graph& graph, const Match& match) { + if (!IsFloat32(*match.valueOf("x")) || !IsFloat32(*match.valueOf("y")) || + !divisor || *divisor == 0.0f || !std::isfinite(1.0f / *divisor)) + return false; ONNX_NAMESPACE::AttributeProto alpha; alpha.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); - alpha.set_f(1.0f / *FloatScalar(graph, *match.valueOf("divisor"))); + alpha.set_f(1.0f / *divisor); match.update("alpha", std::move(alpha)); - return plan; + return true; }; - return Rule{"MatMulDiv", pattern, replacement, guard, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; + const auto* replacement = "(x,y,alpha) => com.microsoft::FusedMatMul(x,y,alpha=alpha)"; + return Rule{"MatMulDiv", pattern, replacement, check_and_bind, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; } ``` -This example uses `IsFloat32` from the first example. The callback is called only -after its guard succeeds. Guards must be pure; neither guard nor callback may -mutate the graph, retain match pointers, or rely on mutable global graph state. -Use local bindings, attributes, and constants to establish applicability. - -`update` accepts a concrete `AttributeProto`, including list-valued attributes. +This example uses `IsFloat32` from the first example. 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 +`Constant` 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. Extra bindings are cleared before each replacement callback; captured attributes -remain available. `attributeOf` retrieves either kind and errors if it is missing. A replacement header +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. +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 @@ -437,7 +448,7 @@ must be connected to their final operator. ## When a match succeeds or fails -Before calling a rule's guard, the matcher checks: +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 @@ -469,7 +480,7 @@ it reports safe rewrite eligibility, not arbitrary structural occurrences. 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 guard check. +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. @@ -487,14 +498,14 @@ 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 guard runs**. The engine does not add +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 guard and +Mixed assignments are rejected. This policy applies before the rule callback and also in match-only mode. To opt into another provider: ```cpp @@ -534,7 +545,7 @@ per call. Rules have deterministic **registration-order priority**, including across different structural groups. At each eligible root, rewriting chooses the first -passing guard. With `Options{.match_only=true}`, all passing guards are counted +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. @@ -576,7 +587,7 @@ 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. Guard false +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 @@ -587,8 +598,8 @@ 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, guard calls, successful matches, and completed -rewrites. Match-only counts all passing guards; rewrite mode counts selected +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 @@ -604,19 +615,19 @@ options.on_failure = [](const MatchFailure& failure) { ``` Stages distinguish opset filtering, structure, type/shape constraints, constants, attributes, deletion -boundaries, execution providers, guards, and replacement validation. An opset +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 guard can use `return match.Reject("unsupported element type");` instead of +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 guard +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. @@ -631,7 +642,7 @@ 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, -guards, boundaries, multi-output rewrites, rule priority, and numerical CPU +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. @@ -646,7 +657,7 @@ 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 semantic guard. Type unions use the separate bounded +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: @@ -656,7 +667,7 @@ 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 guards + 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 @@ -670,7 +681,7 @@ 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 guards; +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. With `indexed=false`, the reference mode instead tries all opset-applicable diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 3f10ca72fc6a3..fd38338d037f0 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -1133,25 +1133,25 @@ class PatternParser { }; } // namespace detail -Rule::Rule(std::string name, Pattern pattern, ReplacementFunction replacement, Guard guard, +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)), - guard(std::move(guard)), + check_and_bind(std::move(check_and_bind)), opsets(std::move(opsets)) {} -Rule::Rule(std::string name, std::string_view pattern, ReplacementFunction replacement, Guard guard, +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(guard), std::move(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, Guard guard, +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(guard), std::move(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, Guard guard, +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(guard), std::move(opsets)) {} + : Rule(std::move(name), Pattern(pattern), Replacement(replacement), std::move(check_and_bind), std::move(opsets)) {} Pattern::Pattern(std::string_view text) { ORT_THROW_IF_ERROR(detail::PatternParser(text).Parse(*this)); } @@ -1207,7 +1207,7 @@ const NodeArg* Match::valueOf(std::string_view name) const { return values[value->handle]; } -void Match::update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) const { +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."); @@ -1280,8 +1280,8 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) 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.guard && rule.replacement, - "Rules need a name, at least one node, a guard and a recipe."); + ORT_ENFORCE(!rule.name.empty() && !rule.pattern.Nodes().empty() && rule.check_and_bind && rule.replacement, + "Rules need a name, at least one node, a check_and_bind and a recipe."); InlinedHashSet domains; for (auto& requirement : rule.opsets) { if (requirement.domain == kOnnxDomainAlias) requirement.domain = kOnnxDomain; @@ -1420,19 +1420,20 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { } for (size_t id : rule_ids) { if (!options_.match_only && id >= winner) break; - match.pattern = &rules_[id].pattern; - match.ClearUpdates(); - ++statistics.guards; - if (!rules_[id].guard(graph, match)) { - report(id, root->Index(), FailureStage::Guard, - match.RejectionReason().empty() ? std::string_view("Guard rejected match.") : match.RejectionReason()); + Match rule_match = match; + rule_match.pattern = &rules_[id].pattern; + ++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()); continue; } if (options_.match_only) ++statistics.matches; if (id < winner) { winner = id; - winning_match = match; - winning_pattern = match.pattern; + winning_match = std::move(rule_match); + winning_pattern = winning_match.pattern; } if (!options_.match_only) break; } @@ -1462,7 +1463,6 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { if (it != producers.end()) affected.push_back(it->second); } InlinedVector added; - winning_match.ClearUpdates(); auto recipe = rules_[winner].replacement(graph, winning_match); auto declared = [&](const std::string& node_domain) { const auto& domain = node_domain == kOnnxDomainAlias ? kOnnxDomain : node_domain; diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 4c3af113ee9fc..c66d003203fdd 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -122,18 +122,18 @@ struct Match { InlinedVector nodes; InlinedVector values; const NodeArg* valueOf(std::string_view name) const; - void update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) const; + void update(std::string_view name, ONNX_NAMESPACE::AttributeProto value); const ONNX_NAMESPACE::AttributeProto& attributeOf(std::string_view name) const; const Dimension& dimensionOf(std::string_view name) const; std::optional elementTypeOf(std::string_view name) const; - void ClearUpdates() const { + void ClearUpdates() { extra_attributes_.clear(); rejection_reason_.clear(); } bool BindAttribute(std::string_view name, const ONNX_NAMESPACE::AttributeProto& value); bool BindDimension(std::string_view name, Dimension value); bool BindElementType(std::string_view name, int32_t value); - bool Reject(std::string_view reason) const { + bool Reject(std::string_view reason) { rejection_reason_ = reason; return false; } @@ -144,8 +144,8 @@ struct Match { NodeAttributes captured_attributes_; InlinedHashMap dimensions_; InlinedHashMap element_types_; - mutable NodeAttributes extra_attributes_; - mutable std::string rejection_reason_; + NodeAttributes extra_attributes_; + std::string rejection_reason_; }; struct Replacement { @@ -178,24 +178,24 @@ struct Replacement { }; struct Rule { - using Guard = std::function; + using CheckAndBind = std::function; using ReplacementFunction = std::function; - Rule(std::string name, Pattern pattern, ReplacementFunction replacement, Guard guard, + Rule(std::string name, Pattern pattern, ReplacementFunction replacement, CheckAndBind check_and_bind, InlinedVector opsets); - Rule(std::string name, std::string_view pattern, ReplacementFunction replacement, Guard guard, + 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, Guard guard, + 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, Guard guard, + Rule(std::string name, std::string_view pattern, std::string_view replacement, CheckAndBind check_and_bind, InlinedVector opsets); std::string name; Pattern pattern; - // Guards must be pure. Recipes may set extra attribute bindings, but must not mutate the graph - // or retain pointers into a Match. + // 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; - Guard guard; + CheckAndBind check_and_bind; InlinedVector opsets; }; @@ -205,7 +205,7 @@ struct Statistics { size_t visited = 0; size_t candidates = 0; size_t structural_checks = 0; - size_t guards = 0; + size_t check_and_bind_calls = 0; size_t matches = 0; size_t rewrites = 0; size_t iterations = 0; @@ -219,7 +219,7 @@ enum class FailureStage { Opset, Attribute, Boundary, ExecutionProvider, - Guard, + CheckAndBind, Replacement }; struct MatchFailure { diff --git a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc index 05be0d38a86cb..838d376a749b4 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc @@ -55,9 +55,9 @@ Rules MakeExampleRules() { return Div(mul, 6.0) } )"; - auto guard = [](const Graph&, const Match& m) { return IsFloat32(*m.valueOf("x")); }; + auto check_and_bind = [](const Graph&, Match& m) { return IsFloat32(*m.valueOf("x")); }; const auto* replacement = "(x) => HardSwish(x)"; - Rule rule{"HardSwish", pattern, replacement, guard, {{kOnnxDomain, 14, 21}}}; + Rule rule{"HardSwish", pattern, replacement, check_and_bind, {{kOnnxDomain, 14, 21}}}; rules.push_back(std::move(rule)); } { @@ -67,30 +67,29 @@ Rules MakeExampleRules() { return Mul(sigmoid, x) } )"; - auto guard = [](const Graph&, const Match& m) { return IsFloat32(*m.valueOf("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, guard, {{kOnnxDomain, 14, 21}}}; + 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 guard = [](const Graph& graph, const Match& m) { + auto check_and_bind = [](const Graph& graph, Match& m) { auto scalar = Scalar(graph, *m.valueOf("divisor")); - return IsFloat32(*m.valueOf("x")) && IsFloat32(*m.valueOf("y")) && - scalar && *scalar != 0.0f && std::isfinite(1.0f / *scalar); - }; - Replacement plan("(x, y, alpha) => com.microsoft.FusedMatMul(x, y, alpha=alpha)"); - auto replacement = [plan](const Graph& graph, const Match& m) { - m.update("alpha", FloatAttribute("alpha", 1.0f / *Scalar(graph, *m.valueOf("divisor")))); - return plan; + 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; }; - Rule rule{"MatMulDiv", pattern, replacement, guard, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; + 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 guard = [](const Graph&, const Match& m) { + 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; @@ -108,37 +107,34 @@ Rules MakeExampleRules() { return true; }; const auto* replacement = "(x,y) => com.microsoft.FusedMatMul(x,y,transB=1)"; - Rule rule{"TransposeMatMul", pattern, replacement, guard, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 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 guard = [](const Graph& graph, const Match& m) { + auto check_and_bind = [](const Graph& graph, Match& m) { auto scalar = Scalar(graph, *m.valueOf("divisor")); - auto it = m.nodes[0]->GetAttributes().find("alpha"); - float alpha = it == m.nodes[0]->GetAttributes().end() ? 1.0f : it->second.f(); - return IsFloat32(*m.valueOf("x")) && IsFloat32(*m.valueOf("y")) && - scalar && *scalar != 0.0f && std::isfinite(alpha / *scalar); - }; - Replacement plan(R"( - (x, y, alpha, transA, transB, transBatchA, transBatchB) => - com.microsoft.FusedMatMul(x, y, alpha=alpha, transA=transA, transB=transB, - transBatchA=transBatchA, transBatchB=transBatchB) - )"); - auto replacement = [plan](const Graph& graph, const Match& m) { 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(); - m.update("alpha", FloatAttribute("alpha", alpha / *Scalar(graph, *m.valueOf("divisor")))); + 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 plan; + return true; }; - Rule rule{"FusedMatMulDiv", pattern, replacement, guard, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; + 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)); } { @@ -153,7 +149,7 @@ Rules MakeExampleRules() { return Mul(normalize, scale) } )"; - auto guard = [](const Graph& graph, const Match& m) { + 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(); @@ -167,19 +163,18 @@ Rules MakeExampleRules() { return false; const auto& mean_node = *m.nodes[1]; auto axes = mean_node.GetAttributes().find("axes"); - return 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; + 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; }; - Replacement plan(R"( + const auto* replacement = R"( (x, scale, fused_epsilon) => SimplifiedLayerNormalization(x, scale, axis=-1, epsilon=fused_epsilon, stash_type=1) - )"); - auto replacement = [plan](const Graph& graph, const Match& m) { - m.update("fused_epsilon", FloatAttribute("epsilon", *Scalar(graph, *m.valueOf("epsilon")))); - return plan; - }; - Rule rule{"RmsNormalization", pattern, replacement, guard, {{kOnnxDomain, 7, 17}}}; + )"; + Rule rule{"RmsNormalization", pattern, replacement, check_and_bind, {{kOnnxDomain, 7, 17}}}; rules.push_back(std::move(rule)); } #endif diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index d0ec2cbd5871d..efe81f47ab623 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -76,44 +76,43 @@ size_t Count(const Graph& graph, const std::string& op) { Rule IdentityPair(const char* name = "identity", bool accept = true) { const auto* pattern = "(x) => a = Identity(x); return Identity(a)"; - auto guard = [accept](const Graph&, const Match&) { return accept; }; + auto check_and_bind = [accept](const Graph&, Match&) { return accept; }; const auto* replacement = "(x) => Identity(x)"; - Rule rule{name, pattern, replacement, guard, {{kOnnxDomain, 1, 21}}}; + 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 guard = [](const Graph&, const Match&) { return true; }; - Replacement plan("(x, axis) => { a, b = Split(x, axis=axis); return a, b; }"); - auto replacement = [plan](const Graph&, const Match& match) { + 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 plan; + return true; }; - Rule rule{"IdentitySplit", pattern, replacement, guard, {{kOnnxDomain, 13, 17}}}; + 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 guard = [](const Graph&, const Match& match) { + 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, guard, {{kOnnxDomain, 13, 17}}}; + 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 guard = [](const Graph&, const Match& match) { + 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(); @@ -124,25 +123,24 @@ Rule LayerNormBias() { const auto& scale_width = scale_shape->dim(0); const auto& bias_width = bias_shape->dim(0); auto axis = match.nodes[0]->GetAttributes().find("axis"); - return 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); - }; - Replacement plan(R"( - (x, scale, bias, axis, epsilon, stash_type) => - LayerNormalization(x, scale, bias, axis=axis, epsilon=epsilon, stash_type=stash_type) - )"); - auto replacement = [plan](const Graph&, const Match& match) { + 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 plan; + return true; }; - Rule rule{"LayerNormBias", pattern, replacement, guard, {{kOnnxDomain, 17, 17}}}; + 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; } @@ -332,7 +330,7 @@ TEST(PatternRewriterTests, NamedValuesAndAttributeGuards) { ASSERT_STATUS_OK(test.graph.Resolve()); auto rule = IdentityPair(); rule.pattern = p; - rule.guard = [](const Graph&, const Match& match) { + 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; }; @@ -361,14 +359,14 @@ TEST(PatternRewriterTests, NamedMatchLookupAndReplacementSafety) { ASSERT_STATUS_OK(test.graph.Resolve()); auto first = IdentityPair(), second = IdentityPair(); second.pattern = Pattern("(input) => { middle = Identity(input); root = Identity(middle); return root; }"); - first.guard = [=](const Graph&, const Match& m) { + 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.guard = [=](const Graph&, const Match& m) { + 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); @@ -624,10 +622,10 @@ TEST(PatternRewriterTests, ReplacementPlanBindingAndSequenceInvalidation) { 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.guard = [](const Graph&, const Match&) { return true; }; + second.check_and_bind = [](const Graph&, Match&) { return true; }; second.replacement = Replacement("(second)=>a=Identity(second);return Identity(a)"); auto followup = IdentityPair(); - followup.guard = [](const Graph&, const Match& m) { return m.nodes.back()->Name().find("rewrite") == 0; }; + 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; @@ -669,6 +667,14 @@ TEST(PatternRewriterTests, ReplacementPlanBindingAndSequenceInvalidation) { } 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"); @@ -715,30 +721,25 @@ TEST(PatternRewriterTests, ReplacementExtraAttributeBindings) { ASSERT_STATUS_OK(test.graph.Resolve()); auto first = MakeExampleRules()[0]; auto second = first; - // Deliberate guard misuse verifies that rejected candidates cannot leak updates. - first.guard = [alpha](const Graph&, const Match& m) { + // 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; }; - Replacement plan("(x,alpha,beta)=>h=HardSigmoid(x,alpha=alpha,beta=beta);return Mul(h,x)"); - second.replacement = [plan, scenario, make_attribute](const Graph&, const Match& m) { + 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 plan; // alpha is missing; must not leak from the rejected rule. + 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); - if (scenario == 4) { - return Replacement(R"( - (x,alpha,beta) => - h=HardSigmoid(x,alpha=alpha,beta=beta,extra=0.0) - return Mul(h,x) - )"); - } - return plan; + return true; }; PatternRewriter rewriter({first, second}, {.indexed = indexed}); Statistics stats; @@ -765,14 +766,14 @@ TEST(PatternRewriterTests, ReplacementExtraBindingsVaryPerMatchAndAttributeTypes for (bool indexed : {false, true}) { auto rule = IdentityPair(); rule.pattern = Pattern("(x)=>a=Identity(x);b=Identity(a);return Identity(b)"); - Replacement plan("(x,axis)=>Softmax(x,axis=axis)"); + rule.replacement = Replacement("(x,axis)=>Softmax(x,axis=axis)"); int calls = 0; - rule.replacement = [plan, &calls](const Graph&, const Match& m) { + 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++ == 0 ? 0 : 1); + axis.set_i((calls++ / 2) % 2); m.update("axis", axis); - return plan; + return true; }; PatternRewriter matcher({rule}, {.match_only = true, .indexed = indexed}); PatternRewriter rewriter({rule}, {.indexed = indexed}); @@ -788,10 +789,11 @@ TEST(PatternRewriterTests, ReplacementExtraBindingsVaryPerMatchAndAttributeTypes ASSERT_STATUS_OK(test.graph.Resolve()); Statistics stats; ASSERT_STATUS_OK(matcher.Run(test.graph, stats)); - EXPECT_EQ(calls, repeat); + 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()); @@ -800,16 +802,16 @@ TEST(PatternRewriterTests, ReplacementExtraBindingsVaryPerMatchAndAttributeTypes auto rule = IdentityPair(); rule.pattern = Pattern("(x)=>a=Identity(x);b=Identity(a);return Identity(b)"); - Replacement plan("(x,perm)=>Transpose(x,perm=perm)"); - auto replacement = [plan](const Graph&, const Match& m) { + 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 plan; + return true; }; - rule.replacement = replacement; + rule.check_and_bind = check_and_bind; TransformerTester( [](ModelTestBuilder& b) { auto* x = b.MakeInput({2, 2}, -1.0f, 1.0f); @@ -827,6 +829,97 @@ TEST(PatternRewriterTests, ReplacementExtraBindingsVaryPerMatchAndAttributeTypes 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_float_data(0.0f); + match.update("zero", std::move(zero)); + return true; + }; + Rule rule{"prepared_tensor", + "(x)=>{a=Identity(x);return Identity(a);}", + "(x,zero)=>{z=Constant(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, "Constant"), 1u); + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "Constant") { + const auto& tensor = node.GetAttributes().at("value").t(); + EXPECT_EQ(tensor.data_type(), ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + EXPECT_EQ(tensor.dims_size(), 0); + 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) => @@ -841,7 +934,7 @@ TEST(PatternRewriterTests, ParsedHardSwishConstantsAndExecution) { 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].guard](const Graph& graph, const Match& match) { + 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}}) { @@ -890,7 +983,7 @@ TEST(PatternRewriterTests, ParsedHardSwishConstantsAndExecution) { } auto rule = MakeExampleRules()[0]; rule.pattern = parsed; - rule.guard = semantic_guard; + 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); @@ -899,7 +992,7 @@ TEST(PatternRewriterTests, ParsedHardSwishConstantsAndExecution) { } auto rule = MakeExampleRules()[0]; rule.pattern = parsed; - rule.guard = semantic_guard; + rule.check_and_bind = semantic_guard; TransformerTester([](ModelTestBuilder& b) { BuildHardSwish(b); }, [](InferenceSessionWrapper& session) { EXPECT_EQ(Count(session.GetGraph(), "HardSigmoid"), 1u); @@ -1277,7 +1370,7 @@ TEST(PatternRewriterTests, GenericExecutionProviderPolicy) { if (scenario >= 3) options.compatible_execution_providers.insert(kCudaExecutionProvider); auto rule = IdentityPair(); int guards = 0; - rule.guard = [&guards](const Graph&, const Match&) { + rule.check_and_bind = [&guards](const Graph&, Match&) { ++guards; return true; }; @@ -1327,11 +1420,11 @@ TEST(PatternRewriterTests, PriorityOverlapTerminationAndMatchOnly) { ASSERT_STATUS_OK(test.graph.Resolve()); int first = 0, second = 0; Rule a = IdentityPair("first"), c = IdentityPair("second"); - a.guard = [&first](const Graph&, const Match&) { + a.check_and_bind = [&first](const Graph&, Match&) { ++first; return true; }; - c.guard = [&second](const Graph&, const Match&) { + c.check_and_bind = [&second](const Graph&, Match&) { ++second; return true; }; @@ -1361,7 +1454,7 @@ TEST(PatternRewriterTests, SingleNodeSameSizeAndExpansionNumerics) { 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.guard = [](const Graph& graph, const Match& match) { + numerical_rule.check_and_bind = [](const Graph& graph, Match& match) { const Node* producer = graph.GetProducerNode(match.valueOf("x")->Name()); return !producer || producer->OpType() != "Identity"; }; @@ -1412,7 +1505,7 @@ TEST(PatternRewriterTests, BoundedCyclesPriorityAndMatchOnly) { backward.replacement = Replacement("(x)=>Identity(x)"); auto later = forward; later.name = "later"; - later.guard = [](const Graph&, const Match&) { + later.check_and_bind = [](const Graph&, Match&) { ADD_FAILURE() << "Lower priority rule selected."; return true; }; @@ -1503,7 +1596,7 @@ TEST(PatternRewriterTests, RuleImportRangesAndAliases) { rule.opsets = {{kOnnxDomainAlias, 14, 17}}; rule.pattern = Pattern("(x)=>a=ai.onnx::Identity(x);return ai.onnx::Identity(a)"); int guards = 0; - rule.guard = [&guards](const Graph&, const Match&) { + rule.check_and_bind = [&guards](const Graph&, Match&) { ++guards; return true; }; @@ -1548,15 +1641,15 @@ TEST(PatternRewriterTests, ImportFilteringWithinSharedStructuralGroups) { 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.guard = [&inactive_guards](const Graph&, const Match&) { + inactive.check_and_bind = [&inactive_guards](const Graph&, Match&) { ++inactive_guards; return true; }; - active.guard = [&active_guards](const Graph&, const Match&) { + active.check_and_bind = [&active_guards](const Graph&, Match&) { ++active_guards; return true; }; - later.guard = [&later_guards](const Graph&, const Match&) { + later.check_and_bind = [&later_guards](const Graph&, Match&) { ++later_guards; return true; }; @@ -1587,7 +1680,7 @@ TEST(PatternRewriterTests, MultipleDomainImportRequirements) { auto rule = IdentityPair(); rule.opsets.push_back({kMSDomain, scenario == 1 ? 2 : 1, 2}); int guards = 0; - rule.guard = [&guards](const Graph&, const Match&) { + rule.check_and_bind = [&guards](const Graph&, Match&) { ++guards; return true; }; @@ -1626,20 +1719,20 @@ TEST(PatternRewriterTests, RuleConstructorTextAndParsedForms) { auto check = [](const auto& pattern_text, const auto& replacement_text) { Pattern parsed_pattern(pattern_text); Replacement parsed_replacement(replacement_text); - auto guard = [](const Graph&, const Match&) { return true; }; + 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)", guard, {{kOnnxDomain, 1, 21}}}, - {"text", pattern_text, replacement_text, guard, {{kOnnxDomain, 1, 21}}}, - {"parsed", parsed_pattern, parsed_replacement, guard, {{kOnnxDomain, 1, 21}}}, - {"parsed_pattern", parsed_pattern, replacement_text, guard, {{kOnnxDomain, 1, 21}}}, - {"parsed_replacement", pattern_text, parsed_replacement, guard, {{kOnnxDomain, 1, 21}}}, - {"callback", pattern_text, replacement, guard, {{kOnnxDomain, 1, 21}}}, - {"parsed_callback", parsed_pattern, replacement, guard, {{kOnnxDomain, 1, 21}}}}; + {"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; @@ -1667,8 +1760,8 @@ TEST(PatternRewriterTests, RuleConstructorTextAndParsedForms) { TEST(PatternRewriterTests, RuleConstructorOwnsParsedTextAndReusesPlans) { std::string pattern = "(x)=>{a=Identity(x);return Identity(a);}"; std::string replacement = "(x)=>Identity(x)"; - auto guard = [](const Graph&, const Match&) { return true; }; - Rule rule{"owned", pattern, replacement, guard, {{kOnnxDomain, 1, 21}}}; + 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); @@ -1690,7 +1783,7 @@ TEST(PatternRewriterTests, RuleConstructorOwnsParsedTextAndReusesPlans) { ++calls; return recipe; }; - Rule computed{"computed", rule.pattern, callback, guard, rule.opsets}; + 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); @@ -1702,9 +1795,9 @@ TEST(PatternRewriterTests, RuleConstructorOwnsParsedTextAndReusesPlans) { TEST(PatternRewriterTests, RuleConstructorRejectsMalformedText) { const auto* pattern = "(x)=>{a=Identity(x);return Identity(a);}"; const auto* replacement = "(x)=>Identity(x)"; - auto guard = [](const Graph&, const Match&) { return true; }; + auto check_and_bind = [](const Graph&, Match&) { return true; }; auto make_rule = [&](auto pattern, auto replacement) { - return Rule{"bad", pattern, replacement, guard, {{kOnnxDomain, 1, 21}}}; + 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); @@ -1728,13 +1821,13 @@ TEST(PatternRewriterTests, PriorityAcrossInterleavedStructuralGroups) { ASSERT_STATUS_OK(test.graph.Resolve()); auto rejected = IdentityPair("rejected", false); auto lower_priority = IdentityPair("lower_priority"); - lower_priority.guard = [](const Graph& graph, const Match& m) { + 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 guard = [](const Graph&, const Match&) { return true; }; + auto check_and_bind = [](const Graph&, Match&) { return true; }; const auto* replacement = "(x)=>Identity(x)"; - Rule higher_priority{"higher_priority", pattern, replacement, guard, {{kOnnxDomain, 1, 21}}}; + 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)); @@ -2036,7 +2129,7 @@ TEST(PatternRewriterTests, NamedAttributeCaptures) { auto rule = IdentityPair(); rule.pattern = Pattern("(x)=>{t=Transpose(x,perm=p);return Transpose(t,perm=p);}"); int guards = 0; - rule.guard = [&guards](const Graph&, const Match& match) { + 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); @@ -2115,7 +2208,7 @@ TEST(PatternRewriterTests, OptionalInputSlotsAndNumerics) { ASSERT_STATUS_OK(test.graph.Resolve()); } int guards = 0; - rule.guard = [scenario, &guards](const Graph&, const Match& match) { + 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); @@ -2126,8 +2219,8 @@ TEST(PatternRewriterTests, OptionalInputSlotsAndNumerics) { EXPECT_EQ(stats.rewrites, 1u); EXPECT_EQ(guards, 1); ASSERT_STATUS_OK(test.graph.Resolve()); - // Numerical guard must not retain references to a previous graph's match. - rule.guard = [](const Graph&, const Match&) { return true; }; + // 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) { @@ -2160,7 +2253,7 @@ TEST(PatternRewriterTests, CapturedAttributesVaryAcrossMatches) { test.builder.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); int guards = 0; - rule.guard = [&guards](const Graph&, const Match& match) { + 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); @@ -2174,7 +2267,7 @@ TEST(PatternRewriterTests, CapturedAttributesVaryAcrossMatches) { InlinedHashSet slopes; for (const auto& node : test.graph.Nodes()) slopes.insert(node.GetAttributes().at("alpha").f()); EXPECT_EQ(slopes.size(), 2u); - rule.guard = [](const Graph&, const Match&) { return true; }; + rule.check_and_bind = [](const Graph&, Match&) { return true; }; TransformerTester( build, [](InferenceSessionWrapper& session) { @@ -2231,7 +2324,7 @@ TEST(PatternRewriterTests, StructuredFailureDiagnostics) { b.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); auto rule = IdentityPair("diagnostic"); - if (scenario == 0) rule.guard = [](const Graph&, const Match& match) { return match.Reject("unsupported dtype"); }; + 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) { @@ -2244,7 +2337,7 @@ TEST(PatternRewriterTests, StructuredFailureDiagnostics) { 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::Guard + const auto expected_stage = scenario == 0 ? FailureStage::CheckAndBind : scenario == 1 ? FailureStage::Boundary : scenario == 2 ? FailureStage::ExecutionProvider : scenario == 3 ? FailureStage::Opset @@ -2378,7 +2471,7 @@ TEST(PatternRewriterTests, SharedDimensionAndElementBindings) { 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.guard = [scenario, &guards](const Graph&, const Match& match) { + 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); @@ -2463,7 +2556,7 @@ TEST(PatternRewriterTests, AnnotatedOptionalValuesAndRuleIsolation) { "(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.guard = [](const Graph&, const Match& match) { + 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); @@ -2472,7 +2565,7 @@ TEST(PatternRewriterTests, AnnotatedOptionalValuesAndRuleIsolation) { 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.guard = [](const Graph&, const Match&) { return false; }; + 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);}"); @@ -2542,7 +2635,7 @@ TEST(PatternRewriterTests, TypeUnionRanksAndNone) { ASSERT_STATUS_OK(test.graph.Resolve()); auto rule = IdentityPair(); rule.pattern = Pattern("(x:float[M] | float[M,N])=>return Identity(x)"); - rule.guard = [scenario](const Graph&, const Match& match) { + 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"))); @@ -2568,7 +2661,7 @@ TEST(PatternRewriterTests, TypeUnionRanksAndNone) { ASSERT_STATUS_OK(test.graph.Resolve()); auto rule = IdentityPair(); rule.pattern = Pattern("(x:float[M],low:float | None)=>return Clip(x,low)"); - rule.guard = [scenario](const Graph&, const Match& match) { + rule.check_and_bind = [scenario](const Graph&, Match& match) { EXPECT_EQ(match.valueOf("low") == nullptr, scenario < 2); return true; }; @@ -2595,7 +2688,7 @@ TEST(PatternRewriterTests, TypeUnionGlobalBindingsAndRollback) { : 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.guard = [scenario](const Graph&, const Match& match) { + 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()); @@ -2615,7 +2708,7 @@ TEST(PatternRewriterTests, TypeUnionGlobalBindingsAndRollback) { ASSERT_STATUS_OK(test.graph.Resolve()); auto rule = IdentityPair(); rule.pattern = Pattern("(x,lo:tensor[D] | None)=>return Clip(x,lo)"); - rule.guard = [](const Graph&, const Match& match) { + 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()); @@ -2708,7 +2801,7 @@ TEST(PatternRewriterTests, MultiOutputProducerReachedThroughSecondOutput) { Pattern pattern("(x)=>a,b=Split(x);return Identity(b)"); auto rule = SplitConcat(); rule.pattern = pattern; - rule.guard = [](const Graph&, const Match&) { return true; }; + 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)); @@ -2742,7 +2835,7 @@ TEST(PatternRewriterTests, MultiOutputRootEdgesAndInvalidRecipe) { auto wrong_count = IdentityPair("single_output_split"); wrong_count.pattern = Pattern("(x)=>a=Identity(x);return Split(a)"); auto valid = IdentitySplit(); - valid.guard = [a, b](const Graph&, const Match& match) { + valid.check_and_bind = [a, b](const Graph&, Match& match) { EXPECT_EQ(match.valueOf("a"), a); EXPECT_EQ(match.valueOf("b"), b); return true; @@ -2926,6 +3019,7 @@ 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); @@ -2970,21 +3064,21 @@ TEST(PatternRewriterTests, ExampleGuardsCheckRelevantBindings) { b.SetGraphOutputs(); ASSERT_STATUS_OK(test.graph.Resolve()); auto rule = examples[index]; - auto guard = rule.guard; + auto check_and_bind = rule.check_and_bind; size_t checked = 0; - rule.guard = [&](const Graph& graph, const Match& match) { - EXPECT_TRUE(guard(graph, match)); + 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 guard independently without making the resolved graph ill-typed. + // 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(guard(graph, changed)); + EXPECT_FALSE(check_and_bind(graph, changed)); } } ++checked; - return guard(graph, match); + return check_and_bind(graph, match); }; Statistics stats; ASSERT_STATUS_OK(PatternRewriter({std::move(rule)}).Run(test.graph, stats)); @@ -3009,9 +3103,9 @@ TEST(PatternRewriterTests, BranchedSharedNodeBindings) { 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 guard = [](const Graph&, const Match&) { return true; }; + auto check_and_bind = [](const Graph&, Match&) { return true; }; const auto* replacement = "(x)=>Identity(x)"; - Rule rule{"diamond", pattern, replacement, guard, {{kOnnxDomain, 7, 21}}}; + 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)); From b3c3083c113575969098de1c2452552f047171dc Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Thu, 15 Oct 2026 14:29:10 -0700 Subject: [PATCH 22/45] Keep tensor binding coverage executable and preserve helper preparation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 2 +- .../test/optimizer/pattern_rewriter_test.cc | 16 +++++++++------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index e3638adf80a59..a2e7b71ac5785 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -432,7 +432,7 @@ 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 -`Constant` node's `value` attribute; it is not a new arbitrary tensor-input binding API. +`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 diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index efe81f47ab623..7617faf296dec 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -836,13 +836,14 @@ TEST(PatternRewriterTests, CheckAndBindTensorAttributeToTextReplacement) { 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)=>{a=Identity(x);return Identity(a);}", - "(x,zero)=>{z=Constant(value=zero);return Add(x,z);}", + "(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( @@ -856,12 +857,13 @@ TEST(PatternRewriterTests, CheckAndBindTensorAttributeToTextReplacement) { const auto& graph = session.GetGraph(); EXPECT_EQ(Count(graph, "Identity"), 0u); EXPECT_EQ(Count(graph, "Add"), 1u); - EXPECT_EQ(Count(graph, "Constant"), 1u); + EXPECT_EQ(Count(graph, "ConstantOfShape"), 1u); for (const auto& node : graph.Nodes()) { - if (node.OpType() == "Constant") { + if (node.OpType() == "ConstantOfShape") { const auto& tensor = node.GetAttributes().at("value").t(); EXPECT_EQ(tensor.data_type(), ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - EXPECT_EQ(tensor.dims_size(), 0); + 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); } @@ -2835,10 +2837,10 @@ TEST(PatternRewriterTests, MultiOutputRootEdgesAndInvalidRecipe) { 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](const Graph&, Match& match) { + 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 true; + return prepare(graph, match); }; ASSERT_STATUS_OK(PatternRewriter({wrong_count, valid}, {.indexed = indexed}).Run(test.graph, stats)); EXPECT_EQ(stats.rewrites, 1u); From acb15c11e168239d7f2603fdebd8c8eb679b8d7c Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 16 Oct 2026 10:17:52 -0700 Subject: [PATCH 23/45] Support hash line comments in rewrite definitions Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 14 +++- .../core/optimizer/pattern_rewriter.cc | 10 ++- .../test/optimizer/pattern_rewriter_test.cc | 80 +++++++++++++++++++ 3 files changed, 102 insertions(+), 2 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index a2e7b71ac5785..e41dfb0728818 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -63,7 +63,19 @@ 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. -Nothing may follow the closing brace. Comments are not supported inside definitions. +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"( diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index fd38338d037f0..af46e542f8717 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -763,7 +763,15 @@ class PatternParser { static bool Digit(char c) { return c >= '0' && c <= '9'; } Status Next() { - while (position_ < text_.size() && std::isspace(static_cast(text_[position_]))) ++position_; + 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; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 7617faf296dec..91113b950ace0 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1883,6 +1883,86 @@ TEST(PatternRewriterTests, OutputCountMismatch) { 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); From 9f49711fe1b5d191c25654e1c16fd11adcd68af3 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 16 Oct 2026 10:41:11 -0700 Subject: [PATCH 24/45] Default rewrite checks to acceptance and simplify annotated example Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 32 +++++++----- .../core/optimizer/pattern_rewriter.cc | 18 ++++--- onnxruntime/core/optimizer/pattern_rewriter.h | 18 +++---- .../test/optimizer/pattern_rewriter_test.cc | 52 +++++++++++++++++++ 4 files changed, 90 insertions(+), 30 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index e41dfb0728818..403b0aa5634f4 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -26,32 +26,32 @@ HardSigmoid attributes shown: using namespace onnxruntime; using namespace onnxruntime::rewrite; -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; -} - Rule MakeHardSwishRule() { const auto* pattern = R"( - (x) => { + (x:float[]) => { sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) return Mul(sigmoid, x) } )"; - auto check_and_bind = [](const Graph&, Match& match) { - return IsFloat32(*match.valueOf("x")); - }; const auto* replacement = "(x) => HardSwish(x)"; - return Rule{"HardSigmoidMul", pattern, replacement, check_and_bind, {{kOnnxDomain, 14, 21}}}; + return Rule{"HardSigmoidMul", pattern, replacement, {}, {{kOnnxDomain, 14, 21}}}; } ``` Both occurrences of `x` must bind the **same `NodeArg`**, not just tensors with -equal shapes or contents. The check-and-bind callback checks the relevant data binding `x`; +equal shapes or contents. `float[]` checks that `x` is a float32 tensor of any +rank, including unknown rank (`float` alone would require a scalar); resolved schemas establish the corresponding intermediate and output types. It does not require every matched value to have the same element type. +An empty `check_and_bind` means always accept after the built-in matching and +safety checks. All constructor overloads default it to `{}`. Use the explicit +`{}` placeholder above when supplying the fifth argument, opset requirements. +The opset list defaults to empty, not an inferred version range: if omitted, +populate `rule.opsets` before constructing a `PatternRewriter`. Missing required +domain ranges remain an error. `check_and_bind_calls` counts only callbacks +actually invoked, so an empty callback contributes zero. + These opset bounds are a conservative supported range for this example, not an automatic statement about every schema version. Floating-point rewrites can change evaluation order; validate equivalence within the intended tolerance. @@ -400,6 +400,12 @@ The following rule folds division by a finite float32 scalar into FusedMatMul's #include "core/graph/graph_utils.h" #include "core/optimizer/initializer.h" +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; +} + std::optional FloatScalar(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 || @@ -430,7 +436,7 @@ Rule MakeMatMulDivRule() { } ``` -This example uses `IsFloat32` from the first example. Only the selected rule's +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 diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index af46e542f8717..37a340cebac6b 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -1288,8 +1288,8 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) 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.check_and_bind && rule.replacement, - "Rules need a name, at least one node, a check_and_bind and a recipe."); + 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; @@ -1430,12 +1430,14 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { if (!options_.match_only && id >= winner) break; Match rule_match = match; rule_match.pattern = &rules_[id].pattern; - ++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()); - continue; + 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()); + continue; + } } if (options_.match_only) ++statistics.matches; if (id < winner) { diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index c66d003203fdd..6654f4fbaed1d 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -181,21 +181,21 @@ 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); + 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; + CheckAndBind check_and_bind; // Empty accepts a structurally valid match without additional checks. InlinedVector opsets; }; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 91113b950ace0..cd666f82cbfa7 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -1794,6 +1794,58 @@ TEST(PatternRewriterTests, RuleConstructorOwnsParsedTextAndReusesPlans) { 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)"; From 819f4598b71d760124bf2818c6cc5556eca4b6f7 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 16 Oct 2026 10:55:27 -0700 Subject: [PATCH 25/45] Explain check-and-bind with an attribute-checking example Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/Pattern_Rewriter.md | 83 +++++++++++++++++++++------------------- 1 file changed, 44 insertions(+), 39 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 403b0aa5634f4..298d9fab5d88e 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -44,8 +44,12 @@ rank, including unknown rank (`float` alone would require a scalar); resolved schemas establish the corresponding intermediate and output types. It does not require every matched value to have the same element type. -An empty `check_and_bind` means always accept after the built-in matching and -safety checks. All constructor overloads default it to `{}`. Use the explicit +The fourth `Rule` argument, `check_and_bind`, is an optional C++ callback that +checks semantic conditions not expressed in the pattern and prepares new +match-local bindings for the replacement. It returns `false` to reject a match, +or `true` to accept it with any prepared bindings. Here `{}` means no additional +checks or bindings: accept after the built-in matching and safety checks. +All constructor overloads default it to `{}`. Use the explicit `{}` placeholder above when supplying the fifth argument, opset requirements. The opset list defaults to empty, not an inferred version range: if omitted, populate `rule.opsets` before constructing a `PatternRewriter`. Missing required @@ -391,53 +395,54 @@ 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. -The following rule folds division by a finite float32 scalar into FusedMatMul's -`alpha` attribute. It requires the Microsoft-domain opset to be imported already. +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 -#include -#include "core/graph/graph_utils.h" -#include "core/optimizer/initializer.h" - -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; -} +#include "core/optimizer/pattern_rewriter.h" -std::optional FloatScalar(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(); - return std::isfinite(scalar) ? std::optional(scalar) : std::nullopt; -} +using namespace onnxruntime; +using namespace onnxruntime::rewrite; -Rule MakeMatMulDivRule() { - const auto* pattern = "(x,y,divisor) => { mm = MatMul(x,y); return Div(mm,divisor); }"; - auto check_and_bind = [](const Graph& graph, Match& match) { - auto divisor = FloatScalar(graph, *match.valueOf("divisor")); - if (!IsFloat32(*match.valueOf("x")) || !IsFloat32(*match.valueOf("y")) || - !divisor || *divisor == 0.0f || !std::isfinite(1.0f / *divisor)) +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 alpha; - alpha.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_FLOAT); - alpha.set_f(1.0f / *divisor); - match.update("alpha", std::move(alpha)); + 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,alpha) => com.microsoft::FusedMatMul(x,y,alpha=alpha)"; - return Rule{"MatMulDiv", pattern, replacement, check_and_bind, {{kOnnxDomain, 9, 21}, {kMSDomain, 1, 1}}}; + 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}}}; } ``` -Only the selected rule's -replacement runs, with the bindings prepared by its successful check-and-bind +`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. From 47b6e7e6b27db0859f5f77c8dbccadb348e4065d Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Tue, 20 Oct 2026 17:57:52 -0700 Subject: [PATCH 26/45] Add composable rewriter fragments and scoped match views Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 240 +++++++ .../core/optimizer/pattern_rewriter.cc | 650 +++++++++++++++-- onnxruntime/core/optimizer/pattern_rewriter.h | 57 +- .../test/optimizer/pattern_rewriter_test.cc | 677 ++++++++++++++++++ 4 files changed, 1564 insertions(+), 60 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 298d9fab5d88e..629b90331c14d 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -710,3 +710,243 @@ rule wins even across groups, and lower-priority work is skipped when possible. 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. + +## Composable fragment libraries + +`PatternLibrary` parses an owning library of independent definitions and links +named entry patterns. Definitions may refer forward to other definitions. +`library.GetPattern("Fuse")` returns a copyable `Pattern` usable by every +existing `Rule` constructor. The library and its source text need not outlive +the returned pattern. + +```cpp +PatternLibrary library(R"( + interface GatedActivation(g:float[M,N], u:float[M,N]) -> (product:float[M,N]) + + fragment plain implements GatedActivation { + s = Sigmoid(g) + a = Mul(g,s) + p = Mul(a,u) + return (product=p) + } + + fragment scaled implements GatedActivation: (alpha:float[...]) -> { + scaled = Mul(g,alpha) + s = Sigmoid(scaled) + a = Mul(g,s) + p = Mul(a,u) + return (product=p) + } + + union Forms = plain | scaled + + fragment wrapped implements GatedActivation { + act = Forms(g,u) + return (product=act.product) + } + + pattern Fuse(g,u) { + outer = wrapped(g,u) + return Identity(outer.product) + } +)"); +``` + +An **interface** is an abstract tensor binding contract, not a callable graph. +A **fragment** supplies a graph body implementing that contract. It inherits +interface inputs and returns its interface outputs, either positionally or +with explicit `(port=local_value, ...)` mappings. Named mappings may reorder +ports. Output annotations constrain the mapped values as well as input +annotations; they do not manufacture shape evidence. Identity mappings such +as `return (y=x)` are allowed inside fragments and emit no operator. +A fragment can return independent branches, provided the complete expanded +entry pattern still has the single connected root required by the matcher. + +The colon signature extends a fragment with external inputs discovered by +matching. Here callers supply only `g,u`, while `alpha` binds to the input of +the matched Mul. 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. + +A **union** supplies ordered alternatives with the same named interface. +Its members are concrete fragments, not interfaces, entry patterns or other +unions. A named call `act=Forms(g,u)` and the anonymous call +`act=(plain | scaled)(g,u)` have the same matching behavior. Each occurrence +has an independent selected member, available through `alternative()`; +no user-defined tag name is required. Fragments may call fragments and unions +at arbitrary nesting depths within the expansion bounds. Inline operator or +subgraph alternative spellings are not supported; existing `|` type +annotations are unchanged. + +Calls require one assigned instance name. Use `act.product` for common outputs +and C++ views for concrete internals. A single-output instance additionally +allows `act` as shorthand for its 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. + +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. Libraries contain at most 256 definitions. +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 fragment +calls from ONNX calls without changing standalone `Pattern` parsing. + +Each complete candidate passes the existing structural, type/constant, +deletion-boundary, layer, provider and opset checks 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. + +### Instance binding views and handwritten C++ structs + +`match.instance("outer").instance("act")` returns a read-only `BindingView`. +Its `valueOf`, `attributeOf`, `dimensionOf`, and `elementTypeOf` use names +relative to that occurrence. `definition()` identifies the called definition, +`interfaceName()` its contract, and `alternative()` the selected concrete member +(empty for a concrete call). Fully qualified lookups such as +`match.valueOf("outer.act.alpha")` are equivalent. Missing instances and +unavailable member-local names raise errors; they are not nullable optional +fields. Only the selected member's concrete bindings exist. + +Views are non-owning and valid only during the candidate callback. They must +not be retained across graph mutation. 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")) { + ORT_ENFORCE(v.interfaceName() == "GatedActivation"); + } +}; +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.valueOf("alpha")) {} +}; +struct ActivationView { + GatedView ports; + std::variant alternative; + static std::variant Select(BindingView v) { + if (v.alternative() == "plain") return PlainView(v); + ORT_ENFORCE(v.alternative() == "scaled", "Unexpected activation member."); + return ScaledView(v); + } + explicit ActivationView(BindingView v) : ports(v), alternative(Select(v)) {} +}; +struct FusionView { + GatedView outer; + ActivationView act; + explicit FusionView(const Match& m) + : outer(m.instance("outer")), act(m.instance("outer").instance("act")) {} +}; + +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 MakeFragmentFusion(const PatternLibrary& library) { + auto check = [](const Graph& graph, Match& match) { + FusionView view(match); + if (!CheckGatingShapes(view.outer) || !CheckGatingShapes(view.act.ports)) + return match.Reject("Unsupported 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 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{"FragmentActivation", 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. 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. + +### 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 fragment-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. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 37a340cebac6b..b8d7a7290e66e 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -13,6 +13,7 @@ #include #include "core/graph/graph_utils.h" +#include "core/framework/tensorprotoutils.h" #include "core/optimizer/initializer.h" #include "onnx/checker.h" @@ -59,13 +60,14 @@ std::string Signature(const Pattern& pattern) { using Producers = InlinedHashMap; -bool MatchType(const Pattern& pattern, int handle, const TypeConstraint& constraint, Match& match) { +bool MatchType(const Pattern& pattern, const Graph& graph, int handle, const TypeConstraint& constraint, Match& match) { if (pattern.IsAttribute(handle)) { return match.attributeOf(constraint.value_name).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(); @@ -93,23 +95,28 @@ bool MatchType(const Pattern& pattern, int handle, const TypeConstraint& constra return true; } -bool MatchTypes(const Pattern& pattern, Match& match) { - InlinedVector handles; +bool MatchTypes(const Pattern& pattern, const Graph& graph, Match& match) { + using Constraints = InlinedVector; + InlinedVector> groups; for (const auto& [handle, alternatives] : pattern.TypeConstraints()) { - handles.push_back(handle); - for (const auto& constraint : alternatives) { + groups.emplace_back(handle, &alternatives); + } + for (const auto& [handle, alternatives] : pattern.AdditionalTypeConstraints()) + groups.emplace_back(handle, &alternatives); + std::stable_sort(groups.begin(), groups.end(), [](const auto& a, const auto& b) { return a.first < b.first; }); + 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{}); } } - std::sort(handles.begin(), handles.end()); std::function search = [&](size_t index, Match& state) { - if (index == handles.size()) return true; - const int handle = handles[index]; - for (const auto& constraint : pattern.TypeConstraints().at(handle)) { + if (index == groups.size()) return true; + const auto& [handle, alternatives] = groups[index]; + for (const auto& constraint : *alternatives) { Match candidate = state; - if (MatchType(pattern, handle, constraint, candidate) && search(index + 1, candidate)) { + if (MatchType(pattern, graph, handle, constraint, candidate) && search(index + 1, candidate)) { state = std::move(candidate); return true; } @@ -191,7 +198,7 @@ bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, } } // Input variables must be outside the deletion set, even when aliasing another binding. - if (!MatchTypes(pattern, match)) + if (!MatchTypes(pattern, graph, match)) return fail(FailureStage::TypeShape, "Tensor type/rank/dimension or captured attribute type mismatch."); for (int value = 0; value < static_cast(pattern.ValueCount()); ++value) { if (!pattern.IsInput(value)) continue; @@ -547,6 +554,8 @@ Value Pattern::Constant(Literal constant) { } 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; @@ -605,19 +614,51 @@ bool Pattern::IsInput(Value value) const { } namespace detail { +struct Port { + std::string name; + InlinedVector constraints; +}; + +struct Definition { + enum class Kind { Interface, + Fragment, + Union, + Pattern }; + Kind kind; + std::string name; + std::string interface_name; + std::string signature; + std::string body; + InlinedVector inputs; + InlinedVector outputs; + InlinedVector members; + InlinedVector, 1> plans; + bool compiling = false; +}; + +using ResolveDefinition = std::function; + 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) + : text_(text), resolve_(std::move(resolve)), choices_(choices), definition_(definition) {} + + const auto& Radices() const { return radices_; } + Status Declarations(InlinedVector& definitions); + static Status ValidateTypes(const Pattern& pattern); + 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)); + 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); @@ -650,11 +691,33 @@ class PatternParser { 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")); + if (token_ != ",") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + ORT_RETURN_IF_ERROR(Expect(")")); + ORT_RETURN_IF(ports.size() != definition_->outputs.size(), Error("interface return ports differ")); + 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_values = true; + } else { const size_t start = token_start_; do { std::string name; - ORT_RETURN_IF_ERROR(Identifier(name, names.empty())); + 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_ == ":") { @@ -675,34 +738,40 @@ class PatternParser { names.clear(); returned = true; } else { - ORT_RETURN_IF(std::any_of(names.begin(), names.end(), - [](const auto& name) { return name.find('.') != std::string::npos; }), + 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(), Error("return must identify operator outputs")); + 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), Error("cannot return a leaf input")); + ORT_RETURN_IF(pattern_.IsInput(found->second) && !definition_, Error("cannot return a leaf input")); pattern_.outputs_.push_back(found->second); } returned_values = true; } else { - for (const auto& name : names) - ORT_RETURN_IF(symbols_.count(name) || pattern_.attribute_names_.count(name), + 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(Call(names)); - for (auto& [name, constraint] : annotations) - pattern_.type_constraints_.emplace(symbols_.at(name).handle, std::move(constraint)); + for (auto& [name, constraint] : annotations) { + auto found = symbols_.find(name); + ORT_RETURN_IF(found == symbols_.end() || pattern_.instances_.count(name), + Error("annotate fragment ports in their interface")); + pattern_.type_constraints_.emplace(found->second.handle, std::move(constraint)); + } } if (returned) { - if (!returned_values) pattern_.outputs_ = pattern_.Nodes().back().outputs; + if (!returned_values) pattern_.outputs_ = last_outputs_; 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")); @@ -710,36 +779,15 @@ class PatternParser { } } ORT_RETURN_IF(!returned, Error("missing return")); - InlinedHashSet dimension_names; - InlinedHashSet element_names; - size_t combinations = 1; - for (const auto& [handle, alternatives] : pattern_.type_constraints_) { - ORT_RETURN_IF(combinations > 64 / alternatives.size(), Error("type alternatives exceed 64 combinations")); - combinations *= alternatives.size(); - for (const auto& constraint : alternatives) { - if (pattern_.IsAttribute(handle)) { - ORT_RETURN_IF(alternatives.size() != 1 || !constraint.attribute_type || pattern_.IsOptional(handle), - Error("invalid attribute annotation or optional attribute")); - } else { - ORT_RETURN_IF(!constraint.tensor_type, Error("annotation requires an attribute variable")); - } - if (!constraint.element_variable.empty()) element_names.insert(constraint.element_variable); - for (const auto& dim : constraint.dimensions) - if (const auto* name = std::get_if(&dim)) dimension_names.insert(*name); - } - } - ORT_RETURN_IF(dimension_names.size() + element_names.size() > 256, Error("too many annotation variables")); - for (const auto& name : dimension_names) - ORT_RETURN_IF(symbols_.count(name) || pattern_.attribute_names_.count(name) || element_names.count(name), - Error("dimension name conflicts with another variable")); - for (const auto& name : element_names) - ORT_RETURN_IF(symbols_.count(name) || pattern_.attribute_names_.count(name), - Error("element type name conflicts with another variable")); + 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; - for (Value output : pattern_.Nodes().back().outputs) reachable.insert(output.handle); + 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(), @@ -785,6 +833,8 @@ class PatternParser { 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()) { @@ -849,7 +899,8 @@ class PatternParser { 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.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)); @@ -864,6 +915,10 @@ class PatternParser { 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; @@ -927,6 +982,8 @@ class PatternParser { } 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(); } @@ -978,6 +1035,37 @@ class PatternParser { } Status Call(const InlinedVector& names) { + if (resolve_ && token_ == "(") { + Definition anonymous; + anonymous.kind = Definition::Kind::Union; + 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 union member")); + const auto* definition = resolve_(member); + ORT_RETURN_IF(!definition || definition->kind != Definition::Kind::Fragment, + Error("anonymous alternatives require concrete fragments")); + ORT_RETURN_IF(!anonymous.interface_name.empty() && anonymous.interface_name != definition->interface_name, + Error("incompatible union interfaces")); + anonymous.interface_name = definition->interface_name; + anonymous.inputs = definition->inputs; + anonymous.outputs = definition->outputs; + for (const auto& plan : definition->plans) { + Pattern alternative = *plan; + alternative.instances_[""] = {anonymous.name, anonymous.interface_name, member}; + anonymous.plans.push_back(std::make_shared(std::move(alternative))); + ORT_RETURN_IF(anonymous.plans.size() > 64, Error("graph alternatives exceed 64 concrete plans")); + } + if (token_ != "|") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + ORT_RETURN_IF(anonymous.plans.size() < 2, Error("anonymous union needs at least two alternatives")); + ORT_RETURN_IF_ERROR(Expect(")")); + return FragmentCall(names, anonymous); + } ORT_RETURN_IF(kind_ != Kind::Identifier || token_ == "return", Error("expected operator name")); Operator op{token_, ""}; ORT_RETURN_IF_ERROR(Next()); @@ -995,6 +1083,12 @@ class PatternParser { ORT_RETURN_IF(op.name.empty(), Error("expected qualified operator name")); } } + if (resolve_ && op.domain.empty()) { + if (const auto* definition = resolve_(op.name)) + return FragmentCall(names, *definition); + ORT_RETURN_IF(!ONNX_NAMESPACE::OpSchemaRegistry::Schema(op.name), + Error("unknown fragment or unqualified ONNX operator")); + } ORT_RETURN_IF_ERROR(Expect("(")); InlinedVector inputs; InlinedVector attributes; @@ -1014,7 +1108,8 @@ class PatternParser { ORT_RETURN_IF_ERROR(Next()); if (kind_ == Kind::Identifier) { std::string capture; - ORT_RETURN_IF_ERROR(Identifier(capture)); + ORT_RETURN_IF_ERROR(Identifier(capture, replacement_bindings_ != nullptr)); + ORT_RETURN_IF(pattern_.instances_.count(capture), Error("attribute capture conflicts with a fragment instance")); if (replacement_bindings_) { auto found = symbols_.find(capture); ORT_RETURN_IF(found == symbols_.end() || !pattern_.IsInput(found->second), @@ -1110,6 +1205,7 @@ class PatternParser { 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)); @@ -1138,9 +1234,365 @@ class PatternParser { 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_; + Status Header(InlinedVector& ports, std::string& signature); + Status FragmentCall(const InlinedVector& names, const Definition& definition); +}; + +Status PatternParser::ValidateTypes(const Pattern& pattern) { + 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)) { + ORT_RETURN_IF(alternatives.size() != 1 || !type.attribute_type || type.immutable || pattern.IsOptional(handle), + "Invalid attribute annotation or optional attribute."); + } 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.type_constraints_) ORT_RETURN_IF_ERROR(validate(handle, alternatives)); + for (const auto& [handle, alternatives] : pattern.additional_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::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_ == "fragment") + definition.kind = Definition::Kind::Fragment; + else if (token_ == "union") + definition.kind = Definition::Kind::Union; + else if (token_ == "pattern") + definition.kind = Definition::Kind::Pattern; + else + return Error("expected interface, fragment, union 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::Union) { + 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 union member")); + definition.members.push_back(std::move(member)); + if (token_ != "|") break; + ORT_RETURN_IF_ERROR(Next()); + } while (true); + ORT_RETURN_IF(definition.members.size() < 2 || definition.members.size() > 64, + Error("unions need between 2 and 64 members")); + } else { + if (definition.kind == Definition::Kind::Fragment) { + ORT_RETURN_IF(token_ != "implements", Error("fragment requires an interface")); + 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) { + std::string output_signature; + ORT_RETURN_IF_ERROR(Expect("->")); + ORT_RETURN_IF_ERROR(Header(definition.outputs, output_signature)); + ORT_RETURN_IF(definition.outputs.empty(), Error("interface requires output ports")); + 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}) { + 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.type_constraints_.emplace(value.handle, port.constraints); + } + } + ORT_RETURN_IF_ERROR(ValidateTypes(contract)); + } else { + 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::FragmentCall(const InlinedVector& names, const Definition& definition) { + ORT_RETURN_IF(definition.kind == Definition::Kind::Interface, Error("interfaces are not callable")); + ORT_RETURN_IF(definition.kind == Definition::Kind::Pattern, Error("call a fragment, not an entry pattern")); + ORT_RETURN_IF(names.size() != 1, Error("fragment calls require one named instance")); + const std::string& path = names.front(); + ORT_RETURN_IF(pattern_.instances_.count(path), Error("duplicate fragment instance")); + ORT_RETURN_IF_ERROR(Expect("(")); + InlinedVector inputs; + while (token_ != ")") { + std::string actual; + ORT_RETURN_IF_ERROR(Identifier(actual, true)); + auto found = symbols_.find(actual); + ORT_RETURN_IF(found == symbols_.end(), Error("undefined fragment 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("fragment argument count mismatch")); + ORT_RETURN_IF(definition.plans.empty(), Error("fragment 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]; + ORT_RETURN_IF(pattern_.Nodes().size() + source.Nodes().size() > 64, Error("expanded node limit exceeded")); + InlinedVector map(source.ValueCount(), Value{-1}); + for (size_t i = 0; i < inputs.size(); ++i) { + const Value formal = source.NamedValue(definition.inputs[i].name); + ORT_RETURN_IF(source.IsAttribute(formal), Error("fragment tensor ports cannot be attributes")); + map[formal.handle] = inputs[i]; + if (!source.IsOptional(formal)) pattern_.optional_inputs_.erase(inputs[i].handle); + } + auto scoped = [&](std::string_view name) { return path + "." + std::string(name); }; + 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 = scoped(attribute.capture); + pattern_.nodes_.push_back(std::move(copy)); + } + for (const auto& name : source.attribute_names_) pattern_.attribute_names_.insert(scoped(name)); + for (int handle : source.attribute_inputs_) { + pattern_.attribute_inputs_.insert(map[handle].handle); + attribute_inputs_.insert(map[handle].handle); + } + auto constraints = [&](int handle, const auto& alternatives) { + auto copy = alternatives; + for (auto& type : copy) { + 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_.additional_type_constraints_.emplace_back(map[handle].handle, std::move(copy)); + }; + for (const auto& [handle, alternatives] : source.type_constraints_) constraints(handle, alternatives); + for (const auto& [handle, alternatives] : source.additional_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 (last_outputs_.size() == 1) { + symbols_.emplace(path, last_outputs_.front()); + pattern_.aliases_.emplace(path, last_outputs_.front()); + } + 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::Pattern, "Expected an entry pattern: ", 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; + } + + 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); + return &definition; + } + + void Compile(Definition& definition) { + ORT_ENFORCE(!definition.compiling, "Recursive fragment/union definition: ", definition.name); + if (definition.kind == Definition::Kind::Interface || !definition.plans.empty()) return; + definition.compiling = true; + if (definition.kind == Definition::Kind::Union) { + for (const auto& member : definition.members) { + const auto* concrete = Resolve(member); + ORT_ENFORCE(concrete && concrete->kind == Definition::Kind::Fragment, + "Union members must name concrete fragments: ", member); + ORT_ENFORCE(definition.interface_name.empty() || definition.interface_name == concrete->interface_name, + "Incompatible union interfaces: ", definition.name); + definition.interface_name = concrete->interface_name; + definition.inputs = concrete->inputs; + definition.outputs = concrete->outputs; + for (const auto& plan : concrete->plans) { + Pattern copy = *plan; + copy.instances_[""] = {definition.name, definition.interface_name, member}; + definition.plans.push_back(std::make_shared(std::move(copy))); + ORT_ENFORCE(definition.plans.size() <= 64, "Graph alternatives exceed 64 concrete plans."); + } + } + } else { + if (definition.kind == Definition::Kind::Fragment) { + const auto* contract = Resolve(definition.interface_name); + ORT_ENFORCE(contract && contract->kind == Definition::Kind::Interface, + "Unknown fragment 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; + do { + Pattern plan; + PatternParser parser(text, [&](std::string_view name) { return Resolve(name); }, choices, definition.kind == Definition::Kind::Fragment ? &definition : nullptr); + const auto status = parser.Parse(plan); + ORT_ENFORCE(status.IsOK(), "In definition ", definition.name, ": ", status.ErrorMessage()); + ORT_ENFORCE(!plan.Nodes().empty() || definition.kind == Definition::Kind::Fragment, + "Entry patterns must contain an operator."); + if (definition.kind == Definition::Kind::Fragment) { + ORT_ENFORCE(plan.outputs_.size() == definition.outputs.size(), "Fragment return count differs from interface."); + 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, + "Interface output shadows a different local value: ", port.name); + plan.aliases_.insert_or_assign(port.name, output); + if (!port.constraints.empty()) + plan.additional_type_constraints_.emplace_back(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; + } + + InlinedVector definitions_; + InlinedHashMap by_name_; }; } // 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)), @@ -1206,13 +1658,76 @@ Replacement::Replacement(std::string_view text) { const NodeArg* Match::valueOf(std::string_view name) const { ORT_ENFORCE(pattern, "Named lookup requires a matched pattern."); - auto value = std::find_if(pattern->values_.begin(), pattern->values_.end(), - [&](const Value& candidate) { return candidate.name == name; }); - ORT_ENFORCE(!name.empty() && value != pattern->values_.end(), "Unknown pattern value: ", name); - ORT_ENFORCE(value->handle >= 0 && static_cast(value->handle) < values.size() && - (values[value->handle] || pattern->IsOptional(*value)), + const Value value = pattern->NamedValue(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]; + return values[value.handle]; +} + +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 ONNX_NAMESPACE::AttributeProto& BindingView::attributeOf(std::string_view name) const { + return match_->attributeOf(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 fragment instance: ", path); + return BindingView(*match_, path); +} +std::string_view BindingView::definition() const { + ORT_ENFORCE(match_->pattern && match_->pattern->instances_.count(path_), "No definition at binding path: ", path_); + return match_->pattern->instances_.at(path_).definition; +} +std::string_view BindingView::interfaceName() const { + ORT_ENFORCE(match_->pattern && match_->pattern->instances_.count(path_), "No definition at binding path: ", path_); + return match_->pattern->instances_.at(path_).interface_name; +} +std::string_view BindingView::alternative() const { + ORT_ENFORCE(match_->pattern && match_->pattern->instances_.count(path_), "No definition at binding path: ", path_); + return match_->pattern->instances_.at(path_).alternative; +} + +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) { @@ -1222,6 +1737,8 @@ void Match::update(std::string_view name, ONNX_NAMESPACE::AttributeProto 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); @@ -1283,7 +1800,19 @@ bool Match::BindElementType(std::string_view name, int32_t value) { } PatternRewriter::PatternRewriter(Rules rules, Options options) - : GraphTransformer("PatternRewriter"), rules_(std::move(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) { @@ -1306,6 +1835,9 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) 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); @@ -1343,7 +1875,7 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) // 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.TypeConstraints().empty()) + !rule.pattern.TypeConstraints().empty() || !rule.pattern.AdditionalTypeConstraints().empty()) signature += "#" + std::to_string(i); auto [it, inserted] = signatures.emplace(signature, groups_.size()); if (inserted) { diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 6654f4fbaed1d..0a7d194152891 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -17,7 +17,8 @@ namespace rewrite { namespace detail { class PatternParser; -} +class PatternModule; +} // namespace detail struct Value { std::string name; @@ -62,6 +63,7 @@ struct TypeConstraint { int32_t attribute_type = 0; // Scalar/list attribute interpretation, when applicable. bool tensor_type = true; bool absent = false; + bool immutable = false; }; struct OpsetRequirement { @@ -91,9 +93,13 @@ class Pattern { 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& AdditionalTypeConstraints() const { return additional_type_constraints_; } private: friend class detail::PatternParser; + friend class detail::PatternModule; + friend class PatternLibrary; + friend class BindingView; friend struct Replacement; friend struct Match; friend class PatternRewriter; @@ -115,9 +121,56 @@ class Pattern { InlinedHashSet attribute_names_; InlinedHashMap> type_constraints_; InlinedHashSet attribute_inputs_; + InlinedHashMap aliases_; + InlinedVector>> additional_type_constraints_; + struct Instance { + std::string definition; + std::string interface_name; + std::string alternative; + }; + InlinedHashMap instances_; + InlinedVector, 1> alternatives_; Value AddValue(std::string_view name); }; +// A library owns independent interfaces, fragments, unions and named entry patterns. +class PatternLibrary { + public: + explicit PatternLibrary(std::string_view text); + Pattern GetPattern(std::string_view name) const; + + private: + std::shared_ptr module_; +}; + +struct Match; + +// Non-owning, instance-scoped access. Valid only during this candidate's callback. +class BindingView { + public: + const NodeArg* valueOf(std::string_view name) const; + const ONNX_NAMESPACE::AttributeProto& attributeOf(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 concrete call; the selected member name for a union occurrence. + std::string_view alternative() 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 Match* match_; + std::string path_; +}; + +// 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 { InlinedVector nodes; InlinedVector values; @@ -126,6 +179,8 @@ struct Match { const ONNX_NAMESPACE::AttributeProto& attributeOf(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(); diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index cd666f82cbfa7..8a9ca04cdc92e 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -3491,5 +3491,682 @@ TEST(PatternRewriterTests, NumericalDecoderPrefillAndDecode) { } #endif +#ifndef DISABLE_CONTRIB_OPS +namespace { + +const char* ActivationLibrary = R"( + interface GatedActivation(g:float[M,N], u:float[M,N]) -> (product:float[M,N]) + fragment plain implements GatedActivation { + s = Sigmoid(g) + a = Mul(g,s) + p = Mul(a,u) + return (product=p) + } + fragment scaled implements GatedActivation: (alpha:float[...]) -> { + scaled = Mul(g,alpha) + s = Sigmoid(scaled) + a = Mul(g,s) + p = Mul(a,u) + return (product=p) + } + union Forms = plain | scaled + fragment 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); + ORT_ENFORCE(view.alternative() == "scaled", "Unexpected activation member."); + return ScaledView(view); + } + 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{"FragmentActivation", library.GetPattern(entry), replacement, check, {{kOnnxDomain, 17, 17}, {kMSDomain, 1, 1}}}; +} + +} // namespace + +TEST(PatternRewriterTests, FragmentManualViewsAndSyntheticFusion) { + 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, FragmentNamedAndAnonymousAlternativesAreOccurrenceLocal) { + for (const char* call : {"Forms", "(first | second)"}) { + const std::string text = std::string(R"( + interface Unary(x:tensor[N]) -> (y:tensor[N]) + fragment first implements Unary { local=Identity(x); return (y=local) } + fragment second implements Unary { local=Identity(x); return (y=local) } + union 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, FragmentHygienicDimensionsAttributesAndCommonOutputs) { + PatternLibrary library(R"( + interface Permute(x:float[N,M]) -> (y:float[M,N]) + fragment 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, FragmentIdentityPortsRemapAndInterfaceConstraintsConjoin) { + PatternLibrary library(R"( + interface Pass(x:float[N]) -> (y:float[N]) + fragment 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, FragmentOutputAnnotationsAreEnforcedOnAliases) { + const PatternLibrary library(R"( + interface Pass(x:float[]) -> (y:float[3]) + fragment 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, FragmentUnionTypeAndConstantRejectionRetry) { + const PatternLibrary library(R"( + interface Unary(x:float[...]) -> (y:float[...]) + fragment fixed implements Unary: (a:const float[1]) -> { y=Mul(x,a); return (y=y) } + fragment dynamic implements Unary: (a:float[]) -> { y=Mul(x,a); return (y=y) } + union 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, FragmentDeletionProviderAndLayerBoundaries) { + const PatternLibrary library(R"( + interface Unary(x) -> (y) + fragment 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, FragmentReplacementFailureDoesNotRetry) { + PatternLibrary library(R"( + interface Unary(x) -> (y) + fragment first implements Unary { a=Identity(x); y=Identity(a); return (y=y) } + fragment second implements Unary { a=Identity(x); y=Identity(a); return (y=y) } + union 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, FragmentExtraInputsAreUsableByParsedReplacement) { + PatternLibrary library(R"( + interface Unary(x:float[]) -> (y:float[]) + fragment 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, FragmentDifferentRootsAndHostRecipeSelection) { + // The returned pattern owns its linked plans after both text and library die. + const Pattern pattern = [] { + std::string text = R"( + interface Unary(x) -> (y) + fragment positive implements Unary { y=Relu(x); return (y=y) } + fragment negative implements Unary { y=Neg(x); return (y=y) } + union 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, FragmentCommonMultiOutputsPreserveGraphOutputs) { + PatternLibrary library(R"( + interface SplitPair(x:float[4]) -> (left:float[2],right:float[2]) + fragment 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, FragmentLibraryErrorsAndBounds) { + const char* invalid[] = { + "interface I(x)->(y) pattern P(x){a=I(x);return a.y}", + "fragment f implements Missing {return (y=x)}", + "interface I(x)->(y) fragment f implements I {a=f(x);return (y=a.y)}", + "interface I(x)->(y) fragment f implements I {a=g(x);return (y=a.y)} " + "fragment g implements I {a=f(x);return (y=a.y)}", + "interface I(x)->(y) fragment f implements I {y=Identity(x);return (wrong=y)}", + "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y,y=y)}", + "interface I(x)->(y) fragment 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) fragment f implements I:(unused:float[])->{y=Identity(x);return (y=y)}", + "interface I(x)->(y) fragment f implements I {a=Identity(x);y=Identity(a);return (y=a)}", + "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y)} " + "pattern P(x){a=f(x,x);return a.y}", + "interface I(x)->(y) fragment 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) " + "fragment f implements I {y=Identity(x);return (y=y)} " + "fragment g implements J {y=Identity(x);return (y=y)} union U=f|g", + "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y)} union U=f|missing", + "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y)} union U=f|f", + "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y)} " + "pattern P(x){a=(f|missing)(x);return a.y}", + "interface I(x)->(y) fragment 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) + fragment f implements I {return (y=x)} + fragment g implements I {return (y=x)} + union 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, FragmentLiteralConstantsAreRemapped) { + const PatternLibrary library(R"( + interface Shift(x) -> (y) + fragment 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); + } +} + } // namespace test } // namespace onnxruntime From f9c717ae97aca2e40edbe3f3eb588e18b086ba77 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Wed, 21 Oct 2026 09:24:33 -0700 Subject: [PATCH 27/45] Unify reusable rewriter bodies and alternatives under pattern Make every concrete named pattern callable, defer rooted-entry checks to GetPattern, and preserve nested alternative selection views. Update composition tests, manual C++ examples, and syntax documentation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 123 +++-- .../core/optimizer/pattern_rewriter.cc | 267 ++++++---- onnxruntime/core/optimizer/pattern_rewriter.h | 10 +- .../test/optimizer/pattern_rewriter_test.cc | 456 +++++++++++++++--- 4 files changed, 676 insertions(+), 180 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 629b90331c14d..72cf5f573b0bf 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -657,9 +657,10 @@ normally and are not converted to nonmatches. ## Limits and unsupported forms -Definitions are limited to 64 nodes, 256 values, and 64 inputs/outputs per +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. -There is no operator/subgraph alternation, implicit commutation, variadic matching, +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`. @@ -711,10 +712,12 @@ 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. -## Composable fragment libraries +## Composable pattern libraries -`PatternLibrary` parses an owning library of independent definitions and links -named entry patterns. Definitions may refer forward to other definitions. +`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 every existing `Rule` constructor. The library and its source text need not outlive the returned pattern. @@ -723,14 +726,14 @@ the returned pattern. PatternLibrary library(R"( interface GatedActivation(g:float[M,N], u:float[M,N]) -> (product:float[M,N]) - fragment plain implements GatedActivation { + pattern plain implements GatedActivation { s = Sigmoid(g) a = Mul(g,s) p = Mul(a,u) return (product=p) } - fragment scaled implements GatedActivation: (alpha:float[...]) -> { + pattern scaled implements GatedActivation: (alpha:float[...]) -> { scaled = Mul(g,alpha) s = Sigmoid(scaled) a = Mul(g,s) @@ -738,9 +741,9 @@ PatternLibrary library(R"( return (product=p) } - union Forms = plain | scaled + pattern Forms = plain | scaled - fragment wrapped implements GatedActivation { + pattern wrapped implements GatedActivation { act = Forms(g,u) return (product=act.product) } @@ -753,16 +756,50 @@ PatternLibrary library(R"( ``` An **interface** is an abstract tensor binding contract, not a callable graph. -A **fragment** supplies a graph body implementing that contract. It inherits -interface inputs and returns its interface outputs, either positionally or +A **pattern body** can implement that contract, inheriting interface inputs +and returning its interface outputs, either positionally or with explicit `(port=local_value, ...)` mappings. Named mappings may reorder ports. Output annotations constrain the mapped values as well as input annotations; they do not manufacture shape evidence. Identity mappings such -as `return (y=x)` are allowed inside fragments and emit no operator. -A fragment can return independent branches, provided the complete expanded -entry pattern still has the single connected root required by the matcher. +as `return (y=x)` are allowed inside patterns and emit no operator. +A body without `implements` declares its inputs directly, as in `pattern Fuse(g,u)`. +Its output port names come from explicit `(port=value)` mappings or returned +value names (`return act.product` exposes `product`). A returned unnamed +operator call exposes `output`, or `output0`, `output1`, etc. for multiple +outputs. All named patterns, including patterns suitable for entry use, may +be called directly or through alternatives. + +### Calls versus standalone entries + +Reusable bodies may return identity aliases or independent output branches. +Every 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. A violation +raises a use-time diagnostic naming the requested entry. Thus a library +may contain useful callable patterns that cannot independently serve as rules: -The colon signature extends a fragment with external inputs discovered by +```text +pattern identity(x) { return (y=x) } +pattern branches(x,z) { + left = Identity(x) + right = Neg(z) + return (a=left,b=right) +} +pattern join(x,z) { + p = branches(x,z) + q = identity(p.a) + return Add(q.y,p.b) +} +``` + +`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 and its existing entry validation are unchanged. + +The colon signature extends a pattern with external inputs discovered by matching. Here callers supply only `g,u`, while `alpha` binds to the input of the matched Mul. It remains a tensor, not an automatically decoded float. Additional attribute captures can also be declared in that signature using @@ -770,13 +807,35 @@ the existing attribute-annotation syntax. There is no separate capture language. Both `float[]` and `float[...]` continue to mean **any rank**, not scalar. -A **union** supplies ordered alternatives with the same named interface. -Its members are concrete fragments, not interfaces, entry patterns or other -unions. A named call `act=Forms(g,u)` and the anonymous call +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 `act=Forms(g,u)` and the anonymous call `act=(plain | scaled)(g,u)` have the same matching behavior. Each occurrence -has an independent selected member, available through `alternative()`; -no user-defined tag name is required. Fragments may call fragments and unions -at arbitrary nesting depths within the expansion bounds. Inline operator or +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 Inner = plain | scaled +pattern Outer = Inner | plain +pattern Alias = Outer +pattern nested(g,u) { + act = Alias(g,u) + return Identity(act.product) +} +``` + +For a selection of `scaled`, `act` reports definition `Alias`, member `Outer`; +successive `alternativeView()` calls report `Outer`/`Inner`, `Inner`/`scaled`, +then body `scaled` with no member. Anonymous alternatives use the same +representation and retain their immediate named members too. Patterns may +call other patterns at arbitrary nesting depths within the expansion bounds. +Inline operator or subgraph alternative spellings are not supported; existing `|` type annotations are unchanged. @@ -797,7 +856,7 @@ type-annotation combinations. Libraries contain at most 256 definitions. 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 fragment +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, @@ -814,8 +873,14 @@ so multiple alternatives of one logical rule can contribute matches. `match.instance("outer").instance("act")` returns a read-only `BindingView`. Its `valueOf`, `attributeOf`, `dimensionOf`, and `elementTypeOf` use names relative to that occurrence. `definition()` identifies the called definition, -`interfaceName()` its contract, and `alternative()` the selected concrete member -(empty for a concrete call). Fully qualified lookups such as +`interfaceName()` its contract, and `alternative()` the selected immediate member +(empty for a body call). `alternativeView()` enters that member's metadata; +calling it on a body view 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("outer.act.alpha")` are equivalent. Missing instances and unavailable member-local names raise errors; they are not nullable optional fields. Only the selected member's concrete bindings exist. @@ -848,9 +913,9 @@ struct ActivationView { GatedView ports; std::variant alternative; static std::variant Select(BindingView v) { - if (v.alternative() == "plain") return PlainView(v); + if (v.alternative() == "plain") return PlainView(v.alternativeView()); ORT_ENFORCE(v.alternative() == "scaled", "Unexpected activation member."); - return ScaledView(v); + return ScaledView(v.alternativeView()); } explicit ActivationView(BindingView v) : ports(v), alternative(Select(v)) {} }; @@ -875,7 +940,7 @@ bool CheckGatingShapes(const GatedView& v) { return true; } -Rule MakeFragmentFusion(const PatternLibrary& library) { +Rule MakePatternFusion(const PatternLibrary& library) { auto check = [](const Graph& graph, Match& match) { FusionView view(match); if (!CheckGatingShapes(view.outer) || !CheckGatingShapes(view.act.ports)) @@ -906,7 +971,7 @@ Rule MakeFragmentFusion(const PatternLibrary& library) { return Mul(a,u) } )"); - return Rule{"FragmentActivation", library.GetPattern("Fuse"), recipe, check, + return Rule{"PatternActivation", library.GetPattern("Fuse"), recipe, check, {{kOnnxDomain, 17, 17}, {kMSDomain, 1, 1}}}; } ``` @@ -946,7 +1011,7 @@ 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 fragment-valued parameters, construction-time scalar parameters, +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. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index b8d7a7290e66e..3a0cf173af98a 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -621,9 +621,8 @@ struct Port { struct Definition { enum class Kind { Interface, - Fragment, - Union, - Pattern }; + Body, + Alternatives }; Kind kind; std::string name; std::string interface_name; @@ -638,6 +637,23 @@ struct Definition { 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; +} + class PatternParser { public: explicit PatternParser(std::string_view text, InlinedVector* attributes = nullptr, @@ -649,8 +665,10 @@ class PatternParser { : text_(text), resolve_(std::move(resolve)), choices_(choices), definition_(definition) {} const auto& Radices() const { return radices_; } + const auto& ReturnedNames() const { return returned_names_; } 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()); @@ -702,15 +720,22 @@ class PatternParser { 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(")")); - ORT_RETURN_IF(ports.size() != definition_->outputs.size(), Error("interface return ports differ")); - 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); + 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 { @@ -749,6 +774,9 @@ class PatternParser { 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 { @@ -766,12 +794,16 @@ class PatternParser { for (auto& [name, constraint] : annotations) { auto found = symbols_.find(name); ORT_RETURN_IF(found == symbols_.end() || pattern_.instances_.count(name), - Error("annotate fragment ports in their interface")); + Error("annotate the called pattern's ports, not its instance")); pattern_.type_constraints_.emplace(found->second.handle, std::move(constraint)); } } if (returned) { - if (!returned_values) pattern_.outputs_ = last_outputs_; + 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")); @@ -1037,34 +1069,22 @@ class PatternParser { Status Call(const InlinedVector& names) { if (resolve_ && token_ == "(") { Definition anonymous; - anonymous.kind = Definition::Kind::Union; + 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 union member")); + ORT_RETURN_IF(!members.insert(member).second, Error("duplicate alternative member")); const auto* definition = resolve_(member); - ORT_RETURN_IF(!definition || definition->kind != Definition::Kind::Fragment, - Error("anonymous alternatives require concrete fragments")); - ORT_RETURN_IF(!anonymous.interface_name.empty() && anonymous.interface_name != definition->interface_name, - Error("incompatible union interfaces")); - anonymous.interface_name = definition->interface_name; - anonymous.inputs = definition->inputs; - anonymous.outputs = definition->outputs; - for (const auto& plan : definition->plans) { - Pattern alternative = *plan; - alternative.instances_[""] = {anonymous.name, anonymous.interface_name, member}; - anonymous.plans.push_back(std::make_shared(std::move(alternative))); - ORT_RETURN_IF(anonymous.plans.size() > 64, Error("graph alternatives exceed 64 concrete plans")); - } + 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(anonymous.plans.size() < 2, Error("anonymous union needs at least two alternatives")); ORT_RETURN_IF_ERROR(Expect(")")); - return FragmentCall(names, anonymous); + return PatternCall(names, anonymous); } ORT_RETURN_IF(kind_ != Kind::Identifier || token_ == "return", Error("expected operator name")); Operator op{token_, ""}; @@ -1085,9 +1105,9 @@ class PatternParser { } if (resolve_ && op.domain.empty()) { if (const auto* definition = resolve_(op.name)) - return FragmentCall(names, *definition); + return PatternCall(names, *definition); ORT_RETURN_IF(!ONNX_NAMESPACE::OpSchemaRegistry::Schema(op.name), - Error("unknown fragment or unqualified ONNX operator")); + Error("unknown pattern or unqualified ONNX operator")); } ORT_RETURN_IF_ERROR(Expect("(")); InlinedVector inputs; @@ -1109,7 +1129,7 @@ class PatternParser { 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 fragment instance")); + 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), @@ -1239,10 +1259,32 @@ class PatternParser { const Definition* definition_ = nullptr; InlinedVector radices_; InlinedVector last_outputs_; + InlinedVector returned_names_; Status Header(InlinedVector& ports, std::string& signature); - Status FragmentCall(const InlinedVector& names, const Definition& definition); + Status PatternCall(const InlinedVector& names, const Definition& definition); }; +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) { InlinedHashSet dimensions, elements, names; for (const auto& value : pattern.values_) @@ -1310,33 +1352,28 @@ Status PatternParser::Declarations(InlinedVector& definitions) { Definition definition; if (token_ == "interface") definition.kind = Definition::Kind::Interface; - else if (token_ == "fragment") - definition.kind = Definition::Kind::Fragment; - else if (token_ == "union") - definition.kind = Definition::Kind::Union; else if (token_ == "pattern") - definition.kind = Definition::Kind::Pattern; + definition.kind = Definition::Kind::Body; else - return Error("expected interface, fragment, union or pattern"); + 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::Union) { + 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 union member")); + 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() < 2 || definition.members.size() > 64, - Error("unions need between 2 and 64 members")); + ORT_RETURN_IF(definition.members.size() > 64, Error("alternatives are limited to 64 members")); } else { - if (definition.kind == Definition::Kind::Fragment) { - ORT_RETURN_IF(token_ != "implements", Error("fragment requires an interface")); + if (definition.kind == Definition::Kind::Body && token_ == "implements") { ORT_RETURN_IF_ERROR(Next()); ORT_RETURN_IF_ERROR(Identifier(definition.interface_name)); definition.signature = "()"; @@ -1386,26 +1423,25 @@ Status PatternParser::Declarations(InlinedVector& definitions) { return Status::OK(); } -Status PatternParser::FragmentCall(const InlinedVector& names, const Definition& definition) { +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(definition.kind == Definition::Kind::Pattern, Error("call a fragment, not an entry pattern")); - ORT_RETURN_IF(names.size() != 1, Error("fragment calls require one named instance")); + ORT_RETURN_IF(names.size() != 1, Error("pattern calls require one named instance")); const std::string& path = names.front(); - ORT_RETURN_IF(pattern_.instances_.count(path), Error("duplicate fragment instance")); + 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)); auto found = symbols_.find(actual); - ORT_RETURN_IF(found == symbols_.end(), Error("undefined fragment argument")); + 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("fragment argument count mismatch")); - ORT_RETURN_IF(definition.plans.empty(), Error("fragment has no concrete plans")); + 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; @@ -1413,13 +1449,23 @@ Status PatternParser::FragmentCall(const InlinedVector& names, cons const Pattern& source = *definition.plans[choice]; 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); - ORT_RETURN_IF(source.IsAttribute(formal), Error("fragment tensor ports cannot be attributes")); + if (source.IsAttribute(formal)) { + 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)) 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")); @@ -1443,10 +1489,16 @@ Status PatternParser::FragmentCall(const InlinedVector& names, cons 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 = scoped(attribute.capture); + if (!attribute.capture.empty()) attribute.capture = attribute_name(attribute.capture); pattern_.nodes_.push_back(std::move(copy)); } - for (const auto& name : source.attribute_names_) pattern_.attribute_names_.insert(scoped(name)); + 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); @@ -1454,7 +1506,7 @@ Status PatternParser::FragmentCall(const InlinedVector& names, cons auto constraints = [&](int handle, const auto& alternatives) { auto copy = alternatives; for (auto& type : copy) { - type.value_name = scoped(type.value_name); + 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); @@ -1492,13 +1544,33 @@ class PatternModule { 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::Pattern, "Expected an entry pattern: ", name); + 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; } private: + 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."); + 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."); + 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); + ORT_ENFORCE(reachable.size() == plan.ValueCount(), "Pattern ", name, " cannot be a rule entry: unused inputs."); + } + 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. @@ -1508,31 +1580,20 @@ class PatternModule { } void Compile(Definition& definition) { - ORT_ENFORCE(!definition.compiling, "Recursive fragment/union definition: ", definition.name); + ORT_ENFORCE(!definition.compiling, "Recursive pattern definition: ", definition.name); if (definition.kind == Definition::Kind::Interface || !definition.plans.empty()) return; definition.compiling = true; - if (definition.kind == Definition::Kind::Union) { + if (definition.kind == Definition::Kind::Alternatives) { for (const auto& member : definition.members) { const auto* concrete = Resolve(member); - ORT_ENFORCE(concrete && concrete->kind == Definition::Kind::Fragment, - "Union members must name concrete fragments: ", member); - ORT_ENFORCE(definition.interface_name.empty() || definition.interface_name == concrete->interface_name, - "Incompatible union interfaces: ", definition.name); - definition.interface_name = concrete->interface_name; - definition.inputs = concrete->inputs; - definition.outputs = concrete->outputs; - for (const auto& plan : concrete->plans) { - Pattern copy = *plan; - copy.instances_[""] = {definition.name, definition.interface_name, member}; - definition.plans.push_back(std::make_shared(std::move(copy))); - ORT_ENFORCE(definition.plans.size() <= 64, "Graph alternatives exceed 64 concrete plans."); - } + ORT_ENFORCE(concrete, "Unknown alternative pattern: ", member); + ORT_THROW_IF_ERROR(PatternParser::AppendAlternative(definition, *concrete)); } } else { - if (definition.kind == Definition::Kind::Fragment) { + if (!definition.interface_name.empty()) { const auto* contract = Resolve(definition.interface_name); ORT_ENFORCE(contract && contract->kind == Definition::Kind::Interface, - "Unknown fragment interface: ", definition.interface_name); + "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; @@ -1548,22 +1609,37 @@ class PatternModule { InlinedVector choices; do { Pattern plan; - PatternParser parser(text, [&](std::string_view name) { return Resolve(name); }, choices, definition.kind == Definition::Kind::Fragment ? &definition : nullptr); + PatternParser parser(text, [&](std::string_view name) { return Resolve(name); }, choices, &definition); const auto status = parser.Parse(plan); ORT_ENFORCE(status.IsOK(), "In definition ", definition.name, ": ", status.ErrorMessage()); - ORT_ENFORCE(!plan.Nodes().empty() || definition.kind == Definition::Kind::Fragment, - "Entry patterns must contain an operator."); - if (definition.kind == Definition::Kind::Fragment) { - ORT_ENFORCE(plan.outputs_.size() == definition.outputs.size(), "Fragment return count differs from interface."); + if (definition.interface_name.empty()) { + 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(CompatiblePorts(definition.outputs, outputs), "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, - "Interface output shadows a different local value: ", port.name); + "Output port shadows a different local value in pattern ", definition.name, ": ", port.name); plan.aliases_.insert_or_assign(port.name, output); - if (!port.constraints.empty()) + if (!definition.interface_name.empty() && !port.constraints.empty()) plan.additional_type_constraints_.emplace_back(output.handle, port.constraints); } } @@ -1680,20 +1756,26 @@ std::optional BindingView::elementTypeOf(std::string_view name) const { 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 fragment instance: ", path); + ORT_ENFORCE(match_->pattern->instances_.count(path), "Unknown pattern instance: ", path); return BindingView(*match_, path); } -std::string_view BindingView::definition() const { +const Pattern::Instance& BindingView::Selection() const { ORT_ENFORCE(match_->pattern && match_->pattern->instances_.count(path_), "No definition at binding path: ", path_); - return match_->pattern->instances_.at(path_).definition; + 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::interfaceName() const { - ORT_ENFORCE(match_->pattern && match_->pattern->instances_.count(path_), "No definition at binding path: ", path_); - return match_->pattern->instances_.at(path_).interface_name; -} -std::string_view BindingView::alternative() const { - ORT_ENFORCE(match_->pattern && match_->pattern->instances_.count(path_), "No definition at binding path: ", path_); - return match_->pattern->instances_.at(path_).alternative; +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 { @@ -1742,7 +1824,8 @@ void Match::update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) 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)), + 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)); @@ -1750,6 +1833,10 @@ void Match::update(std::string_view name, ONNX_NAMESPACE::AttributeProto value) } const ONNX_NAMESPACE::AttributeProto& Match::attributeOf(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; auto found = extra_attributes_.find(std::string(name)); diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 0a7d194152891..ddaa8864c7fc1 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -122,18 +122,20 @@ class Pattern { InlinedHashMap> type_constraints_; InlinedHashSet attribute_inputs_; InlinedHashMap aliases_; + InlinedHashMap attribute_aliases_; InlinedVector>> additional_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 independent interfaces, fragments, unions and named entry patterns. +// A library owns interfaces and named patterns (graph bodies or alternatives). class PatternLibrary { public: explicit PatternLibrary(std::string_view text); @@ -155,15 +157,19 @@ class BindingView { BindingView instance(std::string_view name) const; std::string_view definition() const; std::string_view interfaceName() const; - // Empty for a concrete call; the selected member name for a union occurrence. + // 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. diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 8a9ca04cdc92e..863598c42e308 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -3496,21 +3496,21 @@ namespace { const char* ActivationLibrary = R"( interface GatedActivation(g:float[M,N], u:float[M,N]) -> (product:float[M,N]) - fragment plain implements GatedActivation { + pattern plain implements GatedActivation { s = Sigmoid(g) a = Mul(g,s) p = Mul(a,u) return (product=p) } - fragment scaled implements GatedActivation: (alpha:float[...]) -> { + pattern scaled implements GatedActivation: (alpha:float[...]) -> { scaled = Mul(g,alpha) s = Sigmoid(scaled) a = Mul(g,s) p = Mul(a,u) return (product=p) } - union Forms = plain | scaled - fragment wrapped implements GatedActivation { + pattern Forms = plain | scaled + pattern wrapped implements GatedActivation { act = Forms(g,u) return (product=act.product) } @@ -3548,9 +3548,9 @@ struct ActivationView { GatedView ports; std::variant alternative; static std::variant Select(BindingView view) { - if (view.alternative() == "plain") return PlainView(view); + if (view.alternative() == "plain") return PlainView(view.alternativeView()); ORT_ENFORCE(view.alternative() == "scaled", "Unexpected activation member."); - return ScaledView(view); + return ScaledView(view.alternativeView()); } explicit ActivationView(BindingView view) : ports(view), alternative(Select(view)) {} }; @@ -3603,12 +3603,12 @@ Rule ActivationRule(const PatternLibrary& library, std::string_view entry = "Fus return Mul(a,u) } )"); - return Rule{"FragmentActivation", library.GetPattern(entry), replacement, check, {{kOnnxDomain, 17, 17}, {kMSDomain, 1, 1}}}; + return Rule{"PatternCompositionActivation", library.GetPattern(entry), replacement, check, {{kOnnxDomain, 17, 17}, {kMSDomain, 1, 1}}}; } } // namespace -TEST(PatternRewriterTests, FragmentManualViewsAndSyntheticFusion) { +TEST(PatternRewriterTests, PatternCompositionManualViewsAndSyntheticFusion) { const PatternLibrary library(ActivationLibrary); for (bool indexed : {false, true}) { for (int variant = 0; variant < 5; ++variant) { @@ -3679,13 +3679,13 @@ TEST(PatternRewriterTests, FragmentManualViewsAndSyntheticFusion) { } #endif -TEST(PatternRewriterTests, FragmentNamedAndAnonymousAlternativesAreOccurrenceLocal) { +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]) - fragment first implements Unary { local=Identity(x); return (y=local) } - fragment second implements Unary { local=Identity(x); return (y=local) } - union Forms = first | second + 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) } )"; @@ -3734,10 +3734,10 @@ TEST(PatternRewriterTests, FragmentNamedAndAnonymousAlternativesAreOccurrenceLoc } } -TEST(PatternRewriterTests, FragmentHygienicDimensionsAttributesAndCommonOutputs) { +TEST(PatternRewriterTests, PatternCompositionHygienicDimensionsAttributesAndCommonOutputs) { PatternLibrary library(R"( interface Permute(x:float[N,M]) -> (y:float[M,N]) - fragment transpose implements Permute: (axes:ints) -> { + pattern transpose implements Permute: (axes:ints) -> { local=Transpose(x,perm=axes) return (y=local) } @@ -3777,10 +3777,10 @@ TEST(PatternRewriterTests, FragmentHygienicDimensionsAttributesAndCommonOutputs) ASSERT_STATUS_OK(test.graph.Resolve()); } -TEST(PatternRewriterTests, FragmentIdentityPortsRemapAndInterfaceConstraintsConjoin) { +TEST(PatternRewriterTests, PatternCompositionIdentityPortsRemapAndInterfaceConstraintsConjoin) { PatternLibrary library(R"( interface Pass(x:float[N]) -> (y:float[N]) - fragment pass implements Pass { return (y=x) } + 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}) { @@ -3803,10 +3803,10 @@ TEST(PatternRewriterTests, FragmentIdentityPortsRemapAndInterfaceConstraintsConj } } -TEST(PatternRewriterTests, FragmentOutputAnnotationsAreEnforcedOnAliases) { +TEST(PatternRewriterTests, PatternCompositionOutputAnnotationsAreEnforcedOnAliases) { const PatternLibrary library(R"( interface Pass(x:float[]) -> (y:float[3]) - fragment pass implements Pass { return (y=x) } + pattern pass implements Pass { return (y=x) } pattern P(x) { p=pass(x); return Identity(p.y) } )"); for (int64_t size : {2, 3}) { @@ -3824,12 +3824,12 @@ TEST(PatternRewriterTests, FragmentOutputAnnotationsAreEnforcedOnAliases) { } } -TEST(PatternRewriterTests, FragmentUnionTypeAndConstantRejectionRetry) { +TEST(PatternRewriterTests, PatternCompositionUnionTypeAndConstantRejectionRetry) { const PatternLibrary library(R"( interface Unary(x:float[...]) -> (y:float[...]) - fragment fixed implements Unary: (a:const float[1]) -> { y=Mul(x,a); return (y=y) } - fragment dynamic implements Unary: (a:float[]) -> { y=Mul(x,a); return (y=y) } - union Forms = fixed | dynamic + 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; @@ -3855,10 +3855,10 @@ TEST(PatternRewriterTests, FragmentUnionTypeAndConstantRejectionRetry) { EXPECT_EQ(stats.check_and_bind_calls, 1u); } -TEST(PatternRewriterTests, FragmentDeletionProviderAndLayerBoundaries) { +TEST(PatternRewriterTests, PatternCompositionDeletionProviderAndLayerBoundaries) { const PatternLibrary library(R"( interface Unary(x) -> (y) - fragment pair implements Unary { a=Identity(x); y=Identity(a); return (y=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) { @@ -3881,12 +3881,12 @@ TEST(PatternRewriterTests, FragmentDeletionProviderAndLayerBoundaries) { } } -TEST(PatternRewriterTests, FragmentReplacementFailureDoesNotRetry) { +TEST(PatternRewriterTests, PatternCompositionReplacementFailureDoesNotRetry) { PatternLibrary library(R"( interface Unary(x) -> (y) - fragment first implements Unary { a=Identity(x); y=Identity(a); return (y=y) } - fragment second implements Unary { a=Identity(x); y=Identity(a); return (y=y) } - union Forms=first|second + 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; @@ -3909,10 +3909,10 @@ TEST(PatternRewriterTests, FragmentReplacementFailureDoesNotRetry) { ASSERT_STATUS_OK(test.graph.Resolve()); } -TEST(PatternRewriterTests, FragmentExtraInputsAreUsableByParsedReplacement) { +TEST(PatternRewriterTests, PatternCompositionExtraInputsAreUsableByParsedReplacement) { PatternLibrary library(R"( interface Unary(x:float[]) -> (y:float[]) - fragment scale implements Unary: (alpha:float[]) -> { + pattern scale implements Unary: (alpha:float[]) -> { y=Mul(x,alpha) return (y=y) } @@ -3937,14 +3937,14 @@ TEST(PatternRewriterTests, FragmentExtraInputsAreUsableByParsedReplacement) { OnnxRuntimeException); } -TEST(PatternRewriterTests, FragmentDifferentRootsAndHostRecipeSelection) { +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) - fragment positive implements Unary { y=Relu(x); return (y=y) } - fragment negative implements Unary { y=Neg(x); return (y=y) } - union Forms=positive|negative + 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"); @@ -3975,10 +3975,10 @@ TEST(PatternRewriterTests, FragmentDifferentRootsAndHostRecipeSelection) { } } -TEST(PatternRewriterTests, FragmentCommonMultiOutputsPreserveGraphOutputs) { +TEST(PatternRewriterTests, PatternCompositionCommonMultiOutputsPreserveGraphOutputs) { PatternLibrary library(R"( interface SplitPair(x:float[4]) -> (left:float[2],right:float[2]) - fragment split implements SplitPair { + pattern split implements SplitPair { a,b=Split(x,axis=0) return (right=b,left=a) } @@ -4011,33 +4011,33 @@ TEST(PatternRewriterTests, FragmentCommonMultiOutputsPreserveGraphOutputs) { EXPECT_EQ(test.graph.GetOutputs()[1], output); } -TEST(PatternRewriterTests, FragmentLibraryErrorsAndBounds) { +TEST(PatternRewriterTests, PatternCompositionLibraryErrorsAndBounds) { const char* invalid[] = { "interface I(x)->(y) pattern P(x){a=I(x);return a.y}", - "fragment f implements Missing {return (y=x)}", - "interface I(x)->(y) fragment f implements I {a=f(x);return (y=a.y)}", - "interface I(x)->(y) fragment f implements I {a=g(x);return (y=a.y)} " - "fragment g implements I {a=f(x);return (y=a.y)}", - "interface I(x)->(y) fragment f implements I {y=Identity(x);return (wrong=y)}", - "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y,y=y)}", - "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=z)}", + "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) fragment f implements I:(unused:float[])->{y=Identity(x);return (y=y)}", - "interface I(x)->(y) fragment f implements I {a=Identity(x);y=Identity(a);return (y=a)}", - "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y)} " + "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) fragment f implements I {y=Identity(x);return (y=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) " - "fragment f implements I {y=Identity(x);return (y=y)} " - "fragment g implements J {y=Identity(x);return (y=y)} union U=f|g", - "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y)} union U=f|missing", - "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=y)} union U=f|f", - "interface I(x)->(y) fragment f implements I {y=Identity(x);return (y=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) fragment f implements I {y=Identity(x);return (y=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)}", @@ -4052,9 +4052,9 @@ TEST(PatternRewriterTests, FragmentLibraryErrorsAndBounds) { } std::string too_many = R"( interface I(x)->(y) - fragment f implements I {return (y=x)} - fragment g implements I {return (y=x)} - union U=f|g + 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) @@ -4146,10 +4146,10 @@ TEST(PatternRewriterTests, ConstQualifierIsSpecificToItsTypeAlternative) { EXPECT_TRUE(optional.IsOptional(optional.NamedValue("x"))); } -TEST(PatternRewriterTests, FragmentLiteralConstantsAreRemapped) { +TEST(PatternRewriterTests, PatternCompositionLiteralConstantsAreRemapped) { const PatternLibrary library(R"( interface Shift(x) -> (y) - fragment shift implements Shift { y=Add(x,3.0); return (y=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}) { @@ -4168,5 +4168,343 @@ TEST(PatternRewriterTests, FragmentLiteralConstantsAreRemapped) { } } +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); +} + } // namespace test } // namespace onnxruntime From db9e22319e235f56354249dde1e4aed858fc06d1 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Thu, 22 Oct 2026 20:44:11 -0700 Subject: [PATCH 28/45] Fix composable pattern validation and parser portability Preserve tensor interface ports, reject aliased duplicate entry outputs, and bound cumulative expanded calls before scoped metadata copying. Use the portable classic-locale float reader and cover all four review findings with regressions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 14 ++- .../core/optimizer/pattern_rewriter.cc | 22 ++++- onnxruntime/core/optimizer/pattern_rewriter.h | 1 + .../test/optimizer/pattern_rewriter_test.cc | 96 +++++++++++++++++++ 4 files changed, 127 insertions(+), 6 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 72cf5f573b0bf..ae96cdaa80cd6 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -776,8 +776,12 @@ Every 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. A violation -raises a use-time diagnostic naming the requested entry. Thus a library +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. Thus a library may contain useful callable patterns that cannot independently serve as rules: ```text @@ -852,7 +856,11 @@ one annotation remain disjunctive. 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. Libraries contain at most 256 definitions. +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. 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 diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 3a0cf173af98a..a1350eaa46875 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -12,6 +12,7 @@ #include #include +#include "core/common/parse_string.h" #include "core/graph/graph_utils.h" #include "core/framework/tensorprotoutils.h" #include "core/optimizer/initializer.h" @@ -1025,8 +1026,7 @@ class PatternParser { if (!number.empty() && number.front() == '+') number.remove_prefix(1); if (number.find_first_of(".eE") != std::string_view::npos) { float 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() || !std::isfinite(value), + ORT_RETURN_IF(!TryParseStringWithClassicLocale(number, value) || !std::isfinite(value), Error("invalid or nonfinite float literal")); result = value; } else { @@ -1447,6 +1447,10 @@ Status PatternParser::PatternCall(const InlinedVector& names, const 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; @@ -1554,9 +1558,13 @@ class PatternModule { private: 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."); - for (Value output : plan.Outputs()) + 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) { @@ -1612,6 +1620,14 @@ class PatternModule { PatternParser parser(text, [&](std::string_view name) { return Resolve(name); }, choices, &definition); const auto status = parser.Parse(plan); ORT_ENFORCE(status.IsOK(), "In definition ", definition.name, ": ", status.ErrorMessage()); + 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); + for (Value output : plan.outputs_) + ORT_ENFORCE(!plan.IsAttribute(output), + "Interface output must remain a tensor in pattern ", definition.name); + } if (definition.interface_name.empty()) { InlinedVector outputs; InlinedHashSet names; diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index ddaa8864c7fc1..2fe50779e4b15 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -113,6 +113,7 @@ class Pattern { 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_; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 863598c42e308..1116c5009a6c2 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -4506,5 +4506,101 @@ TEST(PatternRewriterTests, UnifiedPatternsExpansionAndAnnotationThresholds) { 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 test } // namespace onnxruntime From 45599ea811eb6bc364d3d8171ced844de1e552a4 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 09:23:11 -0700 Subject: [PATCH 29/45] Preserve composable matcher contracts across grouped dispatch Make structural signatures locale-independent, require mapped interface outputs, preserve global checker order with cached structural matches, and guard contrib-only helpers. Add focused regressions and document the contracts. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 9 +- .../core/optimizer/pattern_rewriter.cc | 84 +++++---- .../core/optimizer/pattern_rewriter_rules.cc | 4 + .../test/optimizer/pattern_rewriter_test.cc | 159 ++++++++++++++++++ 4 files changed, 220 insertions(+), 36 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index ae96cdaa80cd6..959ad374726bd 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -708,6 +708,11 @@ 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. @@ -781,7 +786,9 @@ 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. Thus a 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. Thus a library may contain useful callable patterns that cannot independently serve as rules: ```text diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index a1350eaa46875..8cd0460b39b29 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -26,6 +27,7 @@ std::string Key(const std::string& op, const std::string& domain) { return domai 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) @@ -1627,6 +1629,7 @@ class PatternModule { for (Value output : plan.outputs_) ORT_ENFORCE(!plan.IsAttribute(output), "Interface output must remain a tensor in pattern ", definition.name); + for (Value output : plan.outputs_) plan.optional_inputs_.erase(output.handle); } if (definition.interface_name.empty()) { InlinedVector outputs; @@ -2046,57 +2049,68 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { size_t winner = std::numeric_limits::max(); Match winning_match; const Pattern* winning_pattern = nullptr; - auto consider = [&](const Pattern& pattern, gsl::span rule_ids) { - statistics.candidates += rule_ids.size(); - ++statistics.structural_checks; + struct CandidateMatch { Match match; FailureStage stage = FailureStage::Structure; std::string_view reason; - if (!MatchPattern(pattern, graph, *root, producers, match, stage, reason)) { - for (size_t id : rule_ids) report(id, root->Index(), stage, reason); - return; + 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."; } - if (!CanRewritePartition(match, options_)) { - for (size_t id : rule_ids) - report(id, root->Index(), FailureStage::ExecutionProvider, "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; } - for (size_t id : rule_ids) { - if (!options_.match_only && id >= winner) break; - Match rule_match = match; - rule_match.pattern = &rules_[id].pattern; - 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()); - continue; - } + Match rule_match = candidate.match; + rule_match.pattern = &rules_[id].pattern; + 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_.match_only) break; + } + 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()) { - for (size_t id : it->second) { - const auto& group = groups_[id]; - const auto& ids = applicable_groups[id]; - if (ids.empty() || (!options_.match_only && ids.front() >= winner)) continue; - consider(group.pattern, ids); + 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) { - consider(rules_[id].pattern, gsl::span(&id, 1)); - if (!options_.match_only && winner != std::numeric_limits::max()) break; + 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; diff --git a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc index 838d376a749b4..432f3b697cadf 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc @@ -11,6 +11,7 @@ namespace onnxruntime { namespace rewrite { namespace { +#ifndef DISABLE_CONTRIB_OPS ONNX_NAMESPACE::AttributeProto FloatAttribute(const char* name, float value) { ONNX_NAMESPACE::AttributeProto attr; attr.set_name(name); @@ -18,6 +19,7 @@ ONNX_NAMESPACE::AttributeProto FloatAttribute(const char* name, float value) { attr.set_f(value); return attr; } +#endif bool IsFloat32(const NodeArg& value) { const auto* type = value.TypeAsProto(); @@ -25,6 +27,7 @@ bool IsFloat32(const NodeArg& value) { 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 || @@ -42,6 +45,7 @@ bool Integer(const Node& node, const char* name, int64_t value) { return it != node.GetAttributes().end() && it->second.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_INT && it->second.i() == value; } +#endif } // namespace Rules MakeExampleRules() { diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 1116c5009a6c2..bbfaf73c34686 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -4602,5 +4603,163 @@ TEST(PatternRewriterTests, InterfacePortsCannotBecomeAttributes) { .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(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()); + } + } + } +} + } // namespace test } // namespace onnxruntime From d939120939e141611945655c3c5a58e8b6a621c6 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 10:43:43 -0700 Subject: [PATCH 30/45] Preserve compositional contracts and bounded matching semantics Keep annotation occurrences ordered through argument remapping, preserve tensor interface requirements across callers, materialize replacement schema defaults, and reject dependency nesting beyond 64 definitions before native recursion. Add focused regressions and document the guarantees and bound. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 22 ++- .../core/optimizer/pattern_rewriter.cc | 65 ++++++--- onnxruntime/core/optimizer/pattern_rewriter.h | 5 +- .../test/optimizer/pattern_rewriter_test.cc | 135 ++++++++++++++++++ 4 files changed, 199 insertions(+), 28 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 959ad374726bd..1b0be1a4a0a37 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -211,7 +211,9 @@ Pattern optional("(x,low:float | None) => return Clip(x,low)"); ``` Alternatives are tried left to right, with annotated variables processed in -definition order. Matching finds the first **globally consistent** combination: +definition order. 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. 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 @@ -282,7 +284,9 @@ 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. +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: @@ -788,7 +792,10 @@ 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. Thus a library +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 @@ -845,7 +852,7 @@ For a selection of `scaled`, `act` reports definition `Alias`, member `Outer`; successive `alternativeView()` calls report `Outer`/`Inner`, `Inner`/`scaled`, then body `scaled` with no member. Anonymous alternatives use the same representation and retain their immediate named members too. Patterns may -call other patterns at arbitrary nesting depths within the expansion bounds. +call other patterns within the expansion and dependency-nesting bounds. Inline operator or subgraph alternative spellings are not supported; existing `|` type annotations are unchanged. @@ -867,7 +874,12 @@ 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. +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. 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 diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 8cd0460b39b29..494e23c333e91 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -99,16 +99,9 @@ bool MatchType(const Pattern& pattern, const Graph& graph, int handle, const Typ } bool MatchTypes(const Pattern& pattern, const Graph& graph, Match& match) { - using Constraints = InlinedVector; - InlinedVector> groups; - for (const auto& [handle, alternatives] : pattern.TypeConstraints()) { - groups.emplace_back(handle, &alternatives); - } - for (const auto& [handle, alternatives] : pattern.AdditionalTypeConstraints()) - groups.emplace_back(handle, &alternatives); - std::stable_sort(groups.begin(), groups.end(), [](const auto& a, const auto& b) { return a.first < b.first; }); + const auto& groups = pattern.OrderedTypeConstraints(); for (const auto& [handle, alternatives] : groups) { - for (const auto& constraint : *alternatives) { + 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{}); @@ -117,7 +110,7 @@ bool MatchTypes(const Pattern& pattern, const Graph& graph, Match& match) { std::function search = [&](size_t index, Match& state) { if (index == groups.size()) return true; const auto& [handle, alternatives] = groups[index]; - for (const auto& constraint : *alternatives) { + for (const auto& constraint : alternatives) { Match candidate = state; if (MatchType(pattern, graph, handle, constraint, candidate) && search(index + 1, candidate)) { state = std::move(candidate); @@ -396,6 +389,11 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R 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); @@ -635,6 +633,7 @@ struct Definition { InlinedVector outputs; InlinedVector members; InlinedVector, 1> plans; + size_t dependency_depth = 1; bool compiling = false; }; @@ -689,6 +688,7 @@ class PatternParser { 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; @@ -794,10 +794,14 @@ class PatternParser { } if (!returned_values) { ORT_RETURN_IF_ERROR(Call(names)); - for (auto& [name, constraint] : annotations) { + 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() || pattern_.instances_.count(name), Error("annotate the called pattern's ports, not its instance")); + pattern_.ordered_type_constraints_.emplace_back(found->second.handle, constraint); pattern_.type_constraints_.emplace(found->second.handle, std::move(constraint)); } } @@ -1288,6 +1292,8 @@ Status PatternParser::AppendAlternative(Definition& target, const Definition& me } Status PatternParser::ValidateTypes(const Pattern& pattern) { + for (int handle : pattern.required_tensor_bindings_) + ORT_RETURN_IF(pattern.IsAttribute(handle), "Interface tensor binding cannot become an attribute."); InlinedHashSet dimensions, elements, names; for (const auto& value : pattern.values_) if (!value.name.empty()) names.insert(value.name); @@ -1312,8 +1318,7 @@ Status PatternParser::ValidateTypes(const Pattern& pattern) { } return Status::OK(); }; - for (const auto& [handle, alternatives] : pattern.type_constraints_) ORT_RETURN_IF_ERROR(validate(handle, alternatives)); - for (const auto& [handle, alternatives] : pattern.additional_type_constraints_) + 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) @@ -1406,7 +1411,7 @@ Status PatternParser::Declarations(InlinedVector& definitions) { 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.type_constraints_.emplace(value.handle, port.constraints); + if (!port.constraints.empty()) contract.ordered_type_constraints_.emplace_back(value.handle, port.constraints); } } ORT_RETURN_IF_ERROR(ValidateTypes(contract)); @@ -1509,6 +1514,8 @@ Status PatternParser::PatternCall(const InlinedVector& names, const 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); auto constraints = [&](int handle, const auto& alternatives) { auto copy = alternatives; for (auto& type : copy) { @@ -1517,10 +1524,9 @@ Status PatternParser::PatternCall(const InlinedVector& names, const for (auto& dim : type.dimensions) if (auto* name = std::get_if(&dim)) *name = scoped(*name); } - pattern_.additional_type_constraints_.emplace_back(map[handle].handle, std::move(copy)); + pattern_.ordered_type_constraints_.emplace_back(map[handle].handle, std::move(copy)); }; - for (const auto& [handle, alternatives] : source.type_constraints_) constraints(handle, alternatives); - for (const auto& [handle, alternatives] : source.additional_type_constraints_) constraints(handle, alternatives); + 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(); @@ -1586,12 +1592,21 @@ class PatternModule { 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) { @@ -1623,12 +1638,16 @@ class PatternModule { const auto status = parser.Parse(plan); ORT_ENFORCE(status.IsOK(), "In definition ", definition.name, ": ", status.ErrorMessage()); if (!definition.interface_name.empty()) { - for (const auto& port : definition.inputs) + 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); - for (Value output : plan.outputs_) + 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.interface_name.empty()) { @@ -1659,7 +1678,7 @@ class PatternModule { "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() && !port.constraints.empty()) - plan.additional_type_constraints_.emplace_back(output.handle, port.constraints); + plan.ordered_type_constraints_.emplace_back(output.handle, port.constraints); } } ORT_THROW_IF_ERROR(PatternParser::ValidateTypes(plan)); @@ -1678,10 +1697,14 @@ class PatternModule { } 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 @@ -1981,7 +2004,7 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) // 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.TypeConstraints().empty() || !rule.pattern.AdditionalTypeConstraints().empty()) + !rule.pattern.OrderedTypeConstraints().empty()) signature += "#" + std::to_string(i); auto [it, inserted] = signatures.emplace(signature, groups_.size()); if (inserted) { diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 2fe50779e4b15..58bad8cc57029 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -93,7 +93,7 @@ class Pattern { 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& AdditionalTypeConstraints() const { return additional_type_constraints_; } + const auto& OrderedTypeConstraints() const { return ordered_type_constraints_; } private: friend class detail::PatternParser; @@ -122,9 +122,10 @@ class Pattern { InlinedHashSet attribute_names_; InlinedHashMap> type_constraints_; InlinedHashSet attribute_inputs_; + InlinedHashSet required_tensor_bindings_; InlinedHashMap aliases_; InlinedHashMap attribute_aliases_; - InlinedVector>> additional_type_constraints_; + InlinedVector>> ordered_type_constraints_; struct Instance { std::string definition; std::string interface_name; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index bbfaf73c34686..603c6f88bd482 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -4761,5 +4761,140 @@ TEST(PatternRewriterTests, InterleavedGroupsDispatchCallbacksInRegistrationOrder } } +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("Interface 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(); + } + } + } + } + } +} + } // namespace test } // namespace onnxruntime From be8de02b38b4c970d532299f02f7e5feee29e148 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 10:59:34 -0700 Subject: [PATCH 31/45] Reject late attribute conversion of composed tensor annotations Retain tensor classification for annotations imported from ordinary callees, including identity aliases, while preserving already-classified typed attribute formals. Diagnose conflicts during library construction and cover nested calls and classification order in both matcher modes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 6 ++ .../core/optimizer/pattern_rewriter.cc | 5 ++ onnxruntime/core/optimizer/pattern_rewriter.h | 1 + .../test/optimizer/pattern_rewriter_test.cc | 57 +++++++++++++++++++ 4 files changed, 69 insertions(+) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 1b0be1a4a0a37..3a34f32799482 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -866,6 +866,12 @@ 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. +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). diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 494e23c333e91..1640174dcc561 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -1294,6 +1294,8 @@ Status PatternParser::AppendAlternative(Definition& target, const Definition& me Status PatternParser::ValidateTypes(const Pattern& pattern) { for (int handle : pattern.required_tensor_bindings_) ORT_RETURN_IF(pattern.IsAttribute(handle), "Interface 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); @@ -1516,7 +1518,10 @@ Status PatternParser::PatternCall(const InlinedVector& names, const } 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); diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 58bad8cc57029..2a9611730b2db 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -123,6 +123,7 @@ class Pattern { InlinedHashMap> type_constraints_; InlinedHashSet attribute_inputs_; InlinedHashSet required_tensor_bindings_; + InlinedHashSet composed_tensor_annotations_; InlinedHashMap aliases_; InlinedHashMap attribute_aliases_; InlinedVector>> ordered_type_constraints_; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 603c6f88bd482..b592ccb3a3676 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -4896,5 +4896,62 @@ TEST(PatternRewriterTests, ForwardDefinitionNestingHasExplicitBoundary) { } } +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); + } + } +} + } // namespace test } // namespace onnxruntime From d6eb407d3ed9a7a574263d504b6a70f508b3ee59 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 11:30:58 -0700 Subject: [PATCH 32/45] Require known equality for repeated selected dimension names Treat singleton named dimensions as permissive captures and repeated names as equality constraints requiring known matching metadata in every selected occurrence. Count only complete annotation choices, preserve scoped composition and retry isolation, and provide named diagnostics. Add regressions and document semantics; explicitly discard nodiscard Status in the throwing test assertion. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 27 ++- .../core/optimizer/pattern_rewriter.cc | 60 ++++- .../test/optimizer/pattern_rewriter_test.cc | 224 +++++++++++++++++- 3 files changed, 291 insertions(+), 20 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 3a34f32799482..ecb1b62cd76f6 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -163,12 +163,27 @@ 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. Repeated -known bindings must agree exactly: integer 2 differs from symbolic `"batch"`, -and `"batch"` differs from `"other"`. An anonymous unknown dimension leaves -the variable unbound and does not conflict with a later known binding. -Consequently, repeated names do **not** prove runtime equality when occurrences -are unknown; check-and-bind callbacks requiring that proof must reject unknown metadata. +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. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 1640174dcc561..57cdd30abeb30 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -63,7 +63,8 @@ std::string Signature(const Pattern& pattern) { using Producers = InlinedHashMap; -bool MatchType(const Pattern& pattern, const Graph& graph, int handle, const TypeConstraint& constraint, Match& match) { +bool MatchType(const Pattern& pattern, const Graph& graph, int handle, const TypeConstraint& constraint, Match& match, + std::string& reason) { if (pattern.IsAttribute(handle)) { return match.attributeOf(constraint.value_name).type() == constraint.attribute_type; } @@ -92,13 +93,16 @@ bool MatchType(const Pattern& pattern, const Graph& graph, int handle, const Typ 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))) return false; + 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) { +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) { @@ -107,14 +111,45 @@ bool MatchTypes(const Pattern& pattern, const Graph& graph, Match& match) { 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 true; + 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) && search(index + 1, candidate)) { - state = std::move(candidate); - return true; + 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; @@ -123,7 +158,7 @@ bool MatchTypes(const Pattern& pattern, const Graph& graph, Match& match) { } bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, const Producers& producers, - Match& match, FailureStage& stage, std::string_view& reason) { + Match& match, FailureStage& stage, std::string& reason) { auto fail = [&](FailureStage failure_stage, std::string_view message) { stage = failure_stage; reason = message; @@ -194,8 +229,11 @@ bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, } } // Input variables must be outside the deletion set, even when aliasing another binding. - if (!MatchTypes(pattern, graph, match)) - return fail(FailureStage::TypeShape, "Tensor type/rank/dimension or captured attribute type mismatch."); + 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]); @@ -2080,7 +2118,7 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { struct CandidateMatch { Match match; FailureStage stage = FailureStage::Structure; - std::string_view reason; + std::string reason; bool matched = false; }; auto match_group = [&](const Pattern& pattern, CandidateMatch& candidate) { diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index b592ccb3a3676..732aa8f5f9c2e 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -2612,7 +2612,6 @@ TEST(PatternRewriterTests, SharedDimensionAndElementBindings) { 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"); - if (scenario == 5) EXPECT_TRUE(std::holds_alternative(match.dimensionOf("M"))); EXPECT_THROW(match.dimensionOf("missing"), OnnxRuntimeException); EXPECT_THROW(match.elementTypeOf("missing"), OnnxRuntimeException); ONNX_NAMESPACE::AttributeProto attribute; @@ -2624,7 +2623,7 @@ TEST(PatternRewriterTests, SharedDimensionAndElementBindings) { }; Statistics stats; ASSERT_STATUS_OK(PatternRewriter({rule}, {.match_only = true, .indexed = indexed}).Run(test.graph, stats)); - const bool expected = scenario == 0 || scenario == 2 || scenario == 5; + const bool expected = scenario == 0 || scenario == 2; EXPECT_EQ(stats.matches, expected ? 1u : 0u); EXPECT_EQ(guards, expected ? 1 : 0); } @@ -4741,7 +4740,7 @@ TEST(PatternRewriterTests, InterleavedGroupsDispatchCallbacksInRegistrationOrder Statistics stats; PatternRewriter rewriter({first, middle, last}, options); if (match_only && throw_late) { - EXPECT_THROW(rewriter.Run(test.graph, stats), OnnxRuntimeException); + 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); @@ -4953,5 +4952,224 @@ TEST(PatternRewriterTests, ComposedTensorAnnotationsCannotBecomeAttributes) { } } +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()); + } + } +} + } // namespace test } // namespace onnxruntime From e6627eb591c4afd5b14ba0ac0142cba84f673702 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 11:54:29 -0700 Subject: [PATCH 33/45] Keep illustrative rewrite rules local to optimizer tests Relocate the unchanged six-rule example factory and helpers into the test anonymous namespace, remove the production source and public declaration, and clarify documentation usage. Production callers author their own rules. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 8 +- onnxruntime/core/optimizer/pattern_rewriter.h | 3 - .../core/optimizer/pattern_rewriter_rules.cc | 189 ------------------ .../test/optimizer/pattern_rewriter_test.cc | 174 ++++++++++++++++ 4 files changed, 180 insertions(+), 194 deletions(-) delete mode 100644 onnxruntime/core/optimizer/pattern_rewriter_rules.cc diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index ecb1b62cd76f6..732b1a21f4efc 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -510,7 +510,11 @@ it reports safe rewrite eligibility, not arbitrary structural occurrences. ### Types in the supplied examples -`MakeExampleRules()` supplies six rules when contrib operators are enabled: +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 | |---|---| @@ -529,7 +533,7 @@ 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 builds supply only the two HardSwish rules. +Contrib-disabled tests construct only the two HardSwish examples. ## Opset and execution-provider applicability diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 2a9611730b2db..dfb232f4a189e 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -323,8 +323,5 @@ class PatternRewriter final : public GraphTransformer { Options options_; }; -// Float32 examples adapted from genuine ONNXScript rules; EP eligibility is controlled by Options. -Rules MakeExampleRules(); - } // namespace rewrite } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc b/onnxruntime/core/optimizer/pattern_rewriter_rules.cc deleted file mode 100644 index 432f3b697cadf..0000000000000 --- a/onnxruntime/core/optimizer/pattern_rewriter_rules.cc +++ /dev/null @@ -1,189 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. -#include "core/optimizer/pattern_rewriter.h" - -#include -#include - -#include "core/graph/graph_utils.h" -#include "core/optimizer/initializer.h" - -namespace onnxruntime { -namespace rewrite { -namespace { -#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 -} // namespace - -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; -} - -} // namespace rewrite -} // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index 732aa8f5f9c2e..d60260b9d386b 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -21,6 +22,179 @@ 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", From 6dfae4de364ad292a9b4c7ffbd295a1f3785773f Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 12:44:10 -0700 Subject: [PATCH 34/45] Lead pattern rewriter documentation with a concrete example Move construction and ownership details after the opening example and use the explicit float[...] any-rank annotation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 36 +++++++++++++++++++----------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 732b1a21f4efc..ad0cb442704a5 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -1,20 +1,5 @@ # Writing C++ graph rewrite 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. - -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.** - ## Example: HardSigmoid followed by Mul This complete rule recognizes the HardSwish decomposition with the exact @@ -28,7 +13,7 @@ using namespace onnxruntime::rewrite; Rule MakeHardSwishRule() { const auto* pattern = R"( - (x:float[]) => { + (x:float[...]) => { sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) return Mul(sigmoid, x) } @@ -39,7 +24,7 @@ Rule MakeHardSwishRule() { ``` Both occurrences of `x` must bind the **same `NodeArg`**, not just tensors with -equal shapes or contents. `float[]` checks that `x` is a float32 tensor of any +equal shapes or contents. `float[...]` checks that `x` is a float32 tensor of any rank, including unknown rank (`float` alone would require a scalar); resolved schemas establish the corresponding intermediate and output types. It does not require every matched value to have the same element type. @@ -60,6 +45,23 @@ These opset bounds are a conservative supported range for this example, not an automatic statement about every schema version. Floating-point rewrites can change evaluation order; validate equivalence within the intended tolerance. +## 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. + +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.** + ## Pattern and replacement syntax Definitions start with named inputs and `=>`. Braces around the body are optional; From a9ffc7c4bd7202b02eb28e7083803b7fbdb573ec Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 13:16:41 -0700 Subject: [PATCH 35/45] Introduce composition details after graph alternatives Keep type-alternative discussion local to annotations and move call constraint ordering and graph/type retry distinctions into the composition section. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index ad0cb442704a5..81b6008a9b297 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -227,19 +227,15 @@ Pattern pattern("(x:float[M] | float[M,N]) => return Identity(x)"); Pattern optional("(x,low:float | None) => return Clip(x,low)"); ``` -Alternatives are tried left to right, with annotated variables processed in -definition order. 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. Matching finds the first **globally consistent** combination: +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. -This bounded search applies only to type/shape annotations, not graph structure: -there is still no search over producers or operand permutations. The product -of alternative counts across annotations is limited to 64 at parse time. +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, but `None` is permitted only for header tensor inputs, not assigned outputs or attribute captures. Replacements remain untyped. @@ -886,7 +882,9 @@ 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. +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 @@ -922,6 +920,13 @@ 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. +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. + ### Instance binding views and handwritten C++ structs `match.instance("outer").instance("act")` returns a read-only `BindingView`. From 64739967c6877cab99941dc4a669308443ea4c0f Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 16:31:39 -0700 Subject: [PATCH 36/45] Add positional outputs and expression alternatives to rewrite patterns Support typed positional signatures, explicit multi-output calls, tensor aliases, wildcard actuals, and branch-optional captures while retaining legacy instance-only multi-output calls. Bound nested expansion and document scoped introspection separately from tensor lookup. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 242 +++--- .../core/optimizer/pattern_rewriter.cc | 219 +++++- onnxruntime/core/optimizer/pattern_rewriter.h | 2 + .../test/optimizer/pattern_rewriter_test.cc | 693 +++++++++++++++++- 4 files changed, 1051 insertions(+), 105 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 81b6008a9b297..9f79ab91fdf45 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -750,55 +750,113 @@ the returned pattern. ```cpp PatternLibrary library(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) + 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) - p = Mul(a,u) - return (product=p) + product = Mul(a,u) + return product } - 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) + pattern Fuse(g,u,alpha:float[...] | None) { + product = activation(g,u,alpha) + return Identity(product) } )"); ``` -An **interface** is an abstract tensor binding contract, not a callable graph. -A **pattern body** can implement that contract, inheriting interface inputs -and returning its interface outputs, either positionally or -with explicit `(port=local_value, ...)` mappings. Named mappings may reorder -ports. Output annotations constrain the mapped values as well as input -annotations; they do not manufacture shape evidence. Identity mappings such -as `return (y=x)` are allowed inside patterns and emit no operator. -A body without `implements` declares its inputs directly, as in `pattern Fuse(g,u)`. -Its output port names come from explicit `(port=value)` mappings or returned -value names (`return act.product` exposes `product`). A returned unnamed -operator call exposes `output`, or `output0`, `output1`, etc. for multiple -outputs. All named patterns, including patterns suitable for entry use, may -be called directly or through alternatives. +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 as before. + +`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. + +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 declared input must be used, and every operator must reach at least +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 @@ -816,27 +874,29 @@ Thus a library may contain useful callable patterns that cannot independently serve as rules: ```text -pattern identity(x) { return (y=x) } +pattern identity(x) -> (y) { return x } pattern branches(x,z) { left = Identity(x) right = Neg(z) - return (a=left,b=right) + return left,right } pattern join(x,z) { - p = branches(x,z) - q = identity(p.a) - return Add(q.y,p.b) + 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 and its existing entry validation are unchanged. +syntax remains supported; its operator-root and output-boundary checks are unchanged. -The colon signature extends a pattern with external inputs discovered by -matching. Here callers supply only `g,u`, while `alpha` binds to the input of -the matched Mul. It remains a tensor, not an automatically decoded float. +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 @@ -849,34 +909,40 @@ 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 `act=Forms(g,u)` and the anonymous call -`act=(plain | scaled)(g,u)` have the same matching behavior. Each occurrence +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 Inner = plain | scaled -pattern Outer = Inner | plain +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(g,u) { - act = Alias(g,u) - return Identity(act.product) +pattern nested(x) { + value = Alias(x) + return Identity(value) } ``` -For a selection of `scaled`, `act` reports definition `Alias`, member `Outer`; -successive `alternativeView()` calls report `Outer`/`Inner`, `Inner`/`scaled`, -then body `scaled` with no member. Anonymous alternatives use the same +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. -Inline operator or -subgraph alternative spellings are not supported; existing `|` type -annotations are unchanged. - -Calls require one assigned instance name. Use `act.product` for common outputs -and C++ views for concrete internals. A single-output instance additionally -allows `act` as shorthand for its output. Tensor handles, attribute captures, +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 @@ -905,6 +971,9 @@ 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 @@ -929,20 +998,22 @@ different producers. ### Instance binding views and handwritten C++ structs -`match.instance("outer").instance("act")` returns a read-only `BindingView`. -Its `valueOf`, `attributeOf`, `dimensionOf`, and `elementTypeOf` use names -relative to that occurrence. `definition()` identifies the called definition, +`match.instance("product")` returns a read-only `BindingView` for the activation +call in the example. `view.instance("scaled")` selects its nested expression +occurrence. Its `valueOf`, `optionalValueOf`, `attributeOf`, `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 that member's metadata; -calling it on a body view raises an error. Repeated calls traverse nested +(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("outer.act.alpha")` are equivalent. Missing instances and +`match.valueOf("product.alpha")` are equivalent. Missing instances and unavailable member-local names raise errors; they are not nullable optional -fields. Only the selected member's concrete bindings exist. +fields. Explicitly declared optional tensors are the exception described above. +Only the selected member's concrete bindings exist. Views are non-owning and valid only during the candidate callback. They must not be retained across graph mutation. The following manual structs show what @@ -954,9 +1025,7 @@ struct GatedView { const NodeArg* u; const NodeArg* product; explicit GatedView(BindingView v) - : g(v.valueOf("g")), u(v.valueOf("u")), product(v.valueOf("product")) { - ORT_ENFORCE(v.interfaceName() == "GatedActivation"); - } + : g(v.valueOf("g")), u(v.valueOf("u")), product(v.valueOf("product")) {} }; struct PlainView { GatedView ports; @@ -966,23 +1035,24 @@ struct PlainView { struct ScaledView { GatedView ports; const NodeArg* alpha; - explicit ScaledView(BindingView v) : ports(v), alpha(v.valueOf("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) { - if (v.alternative() == "plain") return PlainView(v.alternativeView()); - ORT_ENFORCE(v.alternative() == "scaled", "Unexpected activation member."); - return ScaledView(v.alternativeView()); + 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 { - GatedView outer; - ActivationView act; - explicit FusionView(const Match& m) - : outer(m.instance("outer")), act(m.instance("outer").instance("act")) {} + ActivationView activation; + explicit FusionView(const Match& m) : activation(m.instance("product")) {} }; bool CheckGatingShapes(const GatedView& v) { @@ -1002,10 +1072,10 @@ bool CheckGatingShapes(const GatedView& v) { Rule MakePatternFusion(const PatternLibrary& library) { auto check = [](const Graph& graph, Match& match) { FusionView view(match); - if (!CheckGatingShapes(view.outer) || !CheckGatingShapes(view.act.ports)) + if (!CheckGatingShapes(view.activation.ports)) return match.Reject("Unsupported gating shapes."); float alpha = 1.0f; - if (const auto* scaled = std::get_if(&view.act.alternative)) { + 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() && @@ -1037,9 +1107,9 @@ Rule MakePatternFusion(const PatternLibrary& library) { 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. This is a small floating-point example, not a +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++; diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 57cdd30abeb30..ec924e6f1a400 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -672,6 +672,7 @@ struct Definition { InlinedVector members; InlinedVector, 1> plans; size_t dependency_depth = 1; + bool explicit_outputs = false; bool compiling = false; }; @@ -701,11 +702,12 @@ class PatternParser { : text_(text), replacement_attributes_(attributes), replacement_bindings_(bindings) {} PatternParser(std::string_view text, ResolveDefinition resolve, gsl::span choices, - const Definition* definition = nullptr) - : text_(text), resolve_(std::move(resolve)), choices_(choices), definition_(definition) {} + 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); @@ -831,14 +833,13 @@ class PatternParser { } } if (!returned_values) { - ORT_RETURN_IF_ERROR(Call(names)); + 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() || pattern_.instances_.count(name), - Error("annotate the called pattern's ports, not its instance")); + 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)); } @@ -874,6 +875,7 @@ class PatternParser { 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(); @@ -1238,7 +1240,9 @@ class PatternParser { } } else { ORT_RETURN_IF(keyword, Error("positional input after attribute")); - if (name == "None") { + if (name == "_" && resolve_) { + inputs.push_back(pattern_.AddValue("")); + } else if (name == "None") { inputs.push_back(Value{-1}); } else { auto found = symbols_.find(name); @@ -1304,8 +1308,14 @@ class PatternParser { 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) { @@ -1331,7 +1341,7 @@ Status PatternParser::AppendAlternative(Definition& target, const Definition& me Status PatternParser::ValidateTypes(const Pattern& pattern) { for (int handle : pattern.required_tensor_bindings_) - ORT_RETURN_IF(pattern.IsAttribute(handle), "Interface tensor binding cannot become an attribute."); + 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; @@ -1388,6 +1398,124 @@ Status PatternParser::Header(InlinedVector& ports, std::string& signature) 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; @@ -1432,11 +1560,13 @@ Status PatternParser::Declarations(InlinedVector& definitions) { } else { ORT_RETURN_IF_ERROR(Header(definition.inputs, definition.signature)); } - if (definition.kind == Definition::Kind::Interface) { + 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)); - ORT_RETURN_IF(definition.outputs.empty(), Error("interface requires output ports")); + 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) { @@ -1455,7 +1585,8 @@ Status PatternParser::Declarations(InlinedVector& definitions) { } } ORT_RETURN_IF_ERROR(ValidateTypes(contract)); - } else { + } + 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()); @@ -1472,7 +1603,9 @@ Status PatternParser::Declarations(InlinedVector& definitions) { 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.size() != 1, Error("pattern calls require one named instance")); + 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("(")); @@ -1480,9 +1613,14 @@ Status PatternParser::PatternCall(const InlinedVector& names, const while (token_ != ")") { std::string actual; ORT_RETURN_IF_ERROR(Identifier(actual, true)); - auto found = symbols_.find(actual); - ORT_RETURN_IF(found == symbols_.end(), Error("undefined pattern argument")); - inputs.push_back(found->second); + 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()); } @@ -1504,13 +1642,14 @@ Status PatternParser::PatternCall(const InlinedVector& names, const 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)) pattern_.optional_inputs_.erase(inputs[i].handle); + 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) { @@ -1579,11 +1718,17 @@ Status PatternParser::PatternCall(const InlinedVector& names, const 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(); } @@ -1627,6 +1772,7 @@ class PatternModule { 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."); } @@ -1675,11 +1821,28 @@ class PatternModule { } 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); + 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); + for (size_t i = 0; i < results.size(); ++i) { + if (it->second[i].constraints.empty()) { + it->second[i] = results[i]; + continue; + } + if (results[i].constraints.empty()) continue; + auto a = it->second[i], b = results[i]; + a.name = b.name = ""; + ORT_ENFORCE(CompatiblePorts({a}, {b}), "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)), @@ -1693,7 +1856,14 @@ class PatternModule { } for (Value output : plan.outputs_) plan.optional_inputs_.erase(output.handle); } - if (definition.interface_name.empty()) { + 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) { @@ -1720,8 +1890,10 @@ class PatternModule { ORT_ENFORCE(local == plan.values_.end() || local->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() && !port.constraints.empty()) + 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)); @@ -1826,11 +1998,21 @@ const NodeArg* Match::valueOf(std::string_view name) const { 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)); } @@ -2040,6 +2222,7 @@ PatternRewriter::PatternRewriter(Rules rules, Options options) 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; diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index dfb232f4a189e..e8b6051fde2a0 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -154,6 +154,7 @@ struct Match; 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; const Dimension& dimensionOf(std::string_view name) const; std::optional elementTypeOf(std::string_view name) const; @@ -184,6 +185,7 @@ struct Match { 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 Dimension& dimensionOf(std::string_view name) const; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index d60260b9d386b..d8757e1e1e24f 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -4994,7 +4994,7 @@ TEST(PatternRewriterTests, ComposedInterfaceTensorBindingsCannotBecomeAttributes 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("Interface tensor binding"), std::string::npos) << error.what(); + EXPECT_NE(std::string(error.what()).find("Declared tensor binding"), std::string::npos) << error.what(); } } } @@ -5345,5 +5345,696 @@ TEST(PatternRewriterTests, DimensionEqualityRejectionRetriesGraphAlternatives) { } } +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, 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 From 90bf0f0fc0604c9d4686bd3f2747103e380db955 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Fri, 23 Oct 2026 16:50:55 -0700 Subject: [PATCH 37/45] Fix positional output reconciliation and binding validation Allow unconstrained inferred result branches without adding annotation occurrences, reject output labels that shadow tensor aliases, and preserve ordinary typed attribute inputs with explicit output signatures. Add direct and nested regression coverage in both matcher modes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 14 +- .../core/optimizer/pattern_rewriter.cc | 35 ++-- .../test/optimizer/pattern_rewriter_test.cc | 164 ++++++++++++++++++ 3 files changed, 201 insertions(+), 12 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 9f79ab91fdf45..eb459366f13d9 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -772,7 +772,19 @@ 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 as before. +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. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index ec924e6f1a400..b7a0146b21fa0 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -695,6 +695,22 @@ bool CompatiblePorts(const InlinedVector& a, const InlinedVector& b) 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, @@ -1577,6 +1593,7 @@ Status PatternParser::Declarations(InlinedVector& definitions) { } 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")); @@ -1831,16 +1848,8 @@ class PatternModule { 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); - for (size_t i = 0; i < results.size(); ++i) { - if (it->second[i].constraints.empty()) { - it->second[i] = results[i]; - continue; - } - if (results[i].constraints.empty()) continue; - auto a = it->second[i], b = results[i]; - a.name = b.name = ""; - ORT_ENFORCE(CompatiblePorts({a}, {b}), "Expression alternative output type mismatch in ", definition.name); - } + ORT_ENFORCE(MergeResultPorts(it->second, results, false), + "Expression alternative output type mismatch in ", definition.name); } } if (!definition.interface_name.empty()) { @@ -1877,7 +1886,8 @@ class PatternModule { if (definition.outputs.empty()) definition.outputs = std::move(outputs); else - ORT_ENFORCE(CompatiblePorts(definition.outputs, outputs), "Pattern alternatives change output ports: ", definition.name); + ORT_ENFORCE(MergeResultPorts(definition.outputs, outputs, true), + "Pattern alternatives change output ports: ", definition.name); } { ORT_ENFORCE(plan.outputs_.size() == definition.outputs.size(), @@ -1889,6 +1899,9 @@ class PatternModule { [&](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); diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index d8757e1e1e24f..e771747f79ee7 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -5962,6 +5962,170 @@ TEST(PatternRewriterTests, ExpressionPlanAndNestingBounds) { } } +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|None)->(y){y=Softmax(x,axis=a);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]) { From 661422fd93c4430ab0db2becc626b3a269e953c2 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Sun, 25 Oct 2026 11:55:22 -0700 Subject: [PATCH 38/45] Expose standalone pattern and replacement components Document component contracts, check-and-bind, and engine resolution boundaries with a two-Transpose example. Add standalone component tests and numerical rule integration coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 306 ++++++++++-- .../core/optimizer/pattern_rewriter.cc | 169 +++++-- onnxruntime/core/optimizer/pattern_rewriter.h | 22 +- .../pattern_rewriter_components_test.cc | 458 ++++++++++++++++++ 4 files changed, 867 insertions(+), 88 deletions(-) create mode 100644 onnxruntime/test/optimizer/pattern_rewriter_components_test.cc diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index eb459366f13d9..29577f55f45ae 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -1,9 +1,13 @@ # Writing C++ graph rewrite rules -## Example: HardSigmoid followed by Mul +## Rewrite rules: composing two Transposes -This complete rule recognizes the HardSwish decomposition with the exact -HardSigmoid attributes shown: +A rule combines a pattern, semantic check-and-bind function, replacement, +and supported opset ranges. This natural example replaces +`Transpose(Transpose(x,p),q)` with one Transpose. 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]]`. ```cpp #include "core/optimizer/pattern_rewriter.h" @@ -11,39 +15,50 @@ HardSigmoid attributes shown: using namespace onnxruntime; using namespace onnxruntime::rewrite; -Rule MakeHardSwishRule() { - const auto* pattern = R"( - (x:float[...]) => { - sigmoid = HardSigmoid(x, alpha=0.16666667, beta=0.5) - return Mul(sigmoid, x) - } - )"; - const auto* replacement = "(x) => HardSwish(x)"; - return Rule{"HardSigmoidMul", pattern, replacement, {}, {{kOnnxDomain, 14, 21}}}; +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& 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; +} + +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}}}; } ``` -Both occurrences of `x` must bind the **same `NodeArg`**, not just tensors with -equal shapes or contents. `float[...]` checks that `x` is a float32 tensor of any -rank, including unknown rank (`float` alone would require a scalar); -resolved schemas establish the corresponding intermediate and output types. -It does not require every matched value to have the same element type. - -The fourth `Rule` argument, `check_and_bind`, is an optional C++ callback that -checks semantic conditions not expressed in the pattern and prepares new -match-local bindings for the replacement. It returns `false` to reject a match, -or `true` to accept it with any prepared bindings. Here `{}` means no additional -checks or bindings: accept after the built-in matching and safety checks. -All constructor overloads default it to `{}`. Use the explicit -`{}` placeholder above when supplying the fifth argument, opset requirements. -The opset list defaults to empty, not an inferred version range: if omitted, -populate `rule.opsets` before constructing a `PatternRewriter`. Missing required -domain ranges remain an error. `check_and_bind_calls` counts only callbacks -actually invoked, so an empty callback contributes zero. - -These opset bounds are a conservative supported range for this example, not an -automatic statement about every schema version. Floating-point rewrites can -change evaluation order; validate equivalence within the intended tolerance. +The pattern requires explicit `perm` attributes; it does not match default +reverse-axis Transposes. 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 rejects an inner +Transpose result exposed as a graph output or consumed elsewhere. The rank-three +noncommuting permutations `[1,0,2]` then `[0,2,1]` compose to `[1,2,0]`; +tests compare actual CPU results before and after rewriting. + +The sections below describe the components independently. The rewriter engine, +which schedules a collection of rules, is separate from all four components. ## Constructing rules @@ -52,6 +67,9 @@ 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, @@ -61,8 +79,46 @@ 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. -## Pattern and replacement syntax +## 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)}"); +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. +})); +``` + +`VisitMatches` anchors at one existing root in an **already resolved** graph. +It visits successful expanded graph plans in source order, each with fresh +bindings. `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. + +The standalone API uses the same structural, attribute, constant, type/shape, +layer and **deletion-boundary** matcher as the engine. It enumerates rewrite-safe +patterns, not arbitrary subgraph occurrences. It deliberately does not apply +rule opset-range or execution-provider policy and does not run a checker, +construct a replacement or traverse other roots/subgraphs. The first consistent +annotation combination is selected within each concrete graph plan; rejecting +that candidate does not search other annotation combinations. + +### 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 @@ -403,7 +459,102 @@ 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. -## Checking applicability and preparing bindings +## 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. + +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`, +`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 `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 +`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 @@ -412,6 +563,25 @@ 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 `VisitMatches`, without applying any rule: + +```cpp +std::optional accepted; +ORT_RETURN_IF_ERROR(pattern.VisitMatches(graph, root_index, [&](Match& candidate) { + if (!ComposeTransposePermutations(graph, candidate)) return true; // Reject: try next. + accepted = candidate; // Keeps prepared "composed" plus selected-plan ownership. + return false; // Accept: stop enumeration. +})); +``` + +Do **not** directly return the checker boolean from a visitor: its traversal +meaning is opposite. 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. @@ -555,7 +725,7 @@ also in match-only mode. To opt into another provider: ```cpp Options options; options.compatible_execution_providers.insert(kCudaExecutionProvider); -PatternRewriter rewriter(Rules{MakeHardSwishRule()}, options); +PatternRewriter rewriter(Rules{MakeComposeTransposeRule()}, options); ``` Allowing a provider is the caller's assertion of kernel and precision @@ -563,12 +733,12 @@ compatibility, not automatic kernel discovery. Clearing the allowed set permits only unassigned nodes. Replacements inherit the root's provider and layer annotation. -## Applying rules, priority, and errors +## The rewriter engine: applying collections of rules Register the transformer through the normal session pipeline: ```cpp -auto transformer = std::make_unique(Rules{MakeHardSwishRule()}); +auto transformer = std::make_unique(Rules{MakeComposeTransposeRule()}); ORT_RETURN_IF_ERROR(session.RegisterGraphTransformer( std::move(transformer), TransformerLevel::Level1)); ``` @@ -576,17 +746,46 @@ ORT_RETURN_IF_ERROR(session.RegisterGraphTransformer( For one already-resolved graph, use `Run` and explicitly resolve afterwards: ```cpp -PatternRewriter rewriter(Rules{MakeHardSwishRule()}); +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 -and performs normal graph resolution. Callers must not concurrently mutate a +`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 @@ -676,6 +875,26 @@ 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 `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 @@ -1027,8 +1246,9 @@ 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 are non-owning and valid only during the candidate callback. They must -not be retained across graph mutation. The following manual structs show what +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 diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index b7a0146b21fa0..6248f5525fdbe 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -289,12 +289,21 @@ bool CanRewritePartition(const Match& match, const Options& options) { [&](const Node* node) { return node->GetExecutionProviderType() == provider; }); } -Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const Replacement& recipe, - Producers& producers, InlinedVector& added) { +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."); - const Node& root = *match.nodes.back(); - ORT_RETURN_IF(recipe.Outputs().size() != pattern.Outputs().size(), + 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); @@ -302,15 +311,21 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R 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] = graph.GetNodeArg(match.values[pattern.Outputs()[i]]->Name()); + 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); - match.attributeOf(binding); + ONNX_NAMESPACE::AttributeProto value; + ORT_RETURN_IF_ERROR(resolve_attribute(binding, 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."); @@ -334,18 +349,9 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R inputs.push_back(nullptr); continue; } - Value value; - if (const auto* handle = std::get_if(&input)) { - value = *handle; - } else { - value = 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."); - const NodeArg* binding = match.values[value]; - inputs.push_back(binding ? graph.GetNodeArg(binding->Name()) : &graph.GetOrCreateNodeArg("", nullptr)); + 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)); } @@ -356,12 +362,10 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R for (const auto& input : nodes[i].inputs) if (const auto* output = std::get_if(&input)) reachable.insert(output->node); } - InlinedVector initializers; - auto rollback = [&]() { - for (auto it = added.rbegin(); it != added.rend(); ++it) graph.RemoveNode(*it); - for (const auto& name : initializers) graph.RemoveInitializedTensor(name); - added.clear(); - }; + 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]; @@ -408,22 +412,14 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R : &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("rewrite_value"), nullptr)); const auto& node = nodes[i]; const auto& domain = node.domain == kOnnxDomainAlias ? kOnnxDomain : node.domain; - NodeAttributes resolved_attributes; - const NodeAttributes* attributes = &node.attributes; - if (!node.attribute_bindings.empty()) { - resolved_attributes = node.attributes; - for (const auto& [name, binding] : node.attribute_bindings) { - auto value = match.attributeOf(binding); - value.set_name(name); - resolved_attributes.emplace(name, std::move(value)); - } - attributes = &resolved_attributes; - } - Node& replacement = graph.AddNode(graph.GenerateNodeName("rewrite"), node.op, "PatternRewriter", inputs, outputs, - root, attributes, 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)) { - rollback(); return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Replacement schema is unavailable: ", node.op); } ORT_TRY { @@ -437,11 +433,34 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R replacement.Op()->Verify(proto); } ORT_CATCH(const ONNX_NAMESPACE::checker::ValidationError& error) { - rollback(); return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid replacement: ", error.what()); } - replacement.SetExecutionProviderType(root.GetExecutionProviderType()); + 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, ONNX_NAMESPACE::AttributeProto& result) -> Status { + result = match.attributeOf(name); + 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) { @@ -1768,7 +1787,6 @@ class PatternModule { return result; } - private: 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; @@ -1793,6 +1811,7 @@ class PatternModule { 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. @@ -1959,6 +1978,71 @@ Rule::Rule(std::string name, std::string_view pattern, std::string_view replacem InlinedVector opsets) : Rule(std::move(name), Pattern(pattern), Replacement(replacement), std::move(check_and_bind), std::move(opsets)) {} +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 { + 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), "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); + } + 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, ONNX_NAMESPACE::AttributeProto& value) -> Status { + 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) { @@ -2334,7 +2418,8 @@ Status PatternRewriter::Run(Graph& graph, Statistics& statistics) const { return; } Match rule_match = candidate.match; - rule_match.pattern = &rules_[id].pattern; + 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)) { diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index e8b6051fde2a0..f61ca6cae1c6b 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -66,6 +66,8 @@ struct TypeConstraint { bool immutable = false; }; +struct Match; + struct OpsetRequirement { std::string domain; int min_version; @@ -94,6 +96,9 @@ class Pattern { 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_; } + // 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; @@ -148,9 +153,7 @@ class PatternLibrary { std::shared_ptr module_; }; -struct Match; - -// Non-owning, instance-scoped access. Valid only during this candidate's callback. +// 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; @@ -182,6 +185,7 @@ Status ReadConstantTensor(const Graph& graph, const NodeArg* value, InlinedVecto 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; @@ -207,6 +211,9 @@ struct Match { 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_; InlinedHashMap dimensions_; InlinedHashMap element_types_; @@ -215,6 +222,11 @@ struct Match { }; struct Replacement { + struct Instantiation { + InlinedVector nodes; + InlinedVector outputs; + }; + using TensorBindings = InlinedHashMap; struct Output { size_t node; size_t slot = 0; @@ -236,6 +248,10 @@ struct Replacement { 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. + Status Instantiate(Graph& graph, const TensorBindings& tensors, const NodeAttributes& attributes, + Instantiation& result, const Node* annotation_source = nullptr) const; private: friend class PatternRewriter; 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..990096e330c4e --- /dev/null +++ b/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc @@ -0,0 +1,458 @@ +// 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, 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 From bcb4daf16da35f512be3cb8962f1682102b886c9 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Mon, 26 Oct 2026 10:01:31 -0700 Subject: [PATCH 39/45] Minor fixes Signed-off-by: Ganesan Ramalingam --- docs/Pattern_Rewriter.md | 65 ++++++++++++++++++++++++---------------- 1 file changed, 40 insertions(+), 25 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 29577f55f45ae..cdf165d60c08c 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -1,13 +1,10 @@ # Writing C++ graph rewrite rules -## Rewrite rules: composing two Transposes +This document describes how to express fusion optimizations as rewrite rules. -A rule combines a pattern, semantic check-and-bind function, replacement, -and supported opset ranges. This natural example replaces -`Transpose(Transpose(x,p),q)` with one Transpose. 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]]`. +## Example: composing two Transposes + +The following example replaces `Transpose(Transpose(x,p),q)` with one Transpose. ```cpp #include "core/optimizer/pattern_rewriter.h" @@ -15,6 +12,38 @@ replacement. For inner permutation `p` and outer `q`, composition is 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 auto* shape = match.valueOf("x")->Shape(); if (!shape) return match.Reject("Transpose composition requires known rank."); @@ -36,29 +65,15 @@ bool ComposeTransposePermutations(const Graph&, Match& match) { match.update("composed", std::move(composed)); return true; } - -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}}}; -} ``` The pattern requires explicit `perm` attributes; it does not match default reverse-axis Transposes. 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 rejects an inner -Transpose result exposed as a graph output or consumed elsewhere. The rank-three -noncommuting permutations `[1,0,2]` then `[0,2,1]` compose to `[1,2,0]`; -tests compare actual CPU results before and after rewriting. - -The sections below describe the components independently. The rewriter engine, -which schedules a collection of rules, is separate from all four components. +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]`. ## Constructing rules From 9570a143c13ce3b1648b8e073c030e24805cb53c Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Mon, 26 Oct 2026 10:17:28 -0700 Subject: [PATCH 40/45] Add direct anchored matching and clarify component documentation Address PR16 feedback with MatchAt first-candidate matching, explicit lookup types, optional attribute constraints, and later rule construction documentation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 103 ++++++++++++------ .../core/optimizer/pattern_rewriter.cc | 10 ++ onnxruntime/core/optimizer/pattern_rewriter.h | 2 + .../pattern_rewriter_components_test.cc | 53 +++++++++ 4 files changed, 134 insertions(+), 34 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index cdf165d60c08c..1a3a666b73390 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -45,11 +45,12 @@ outer `q`, composition is `composed[i]=p[q[i]]`, not `q[p[i]]`. The complete definition of `ComposeTransposePermutations` is: ```cpp bool ComposeTransposePermutations(const Graph&, Match& match) { - const auto* shape = match.valueOf("x")->Shape(); + 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 auto& p = match.attributeOf("p"); - const auto& q = match.attributeOf("q"); + 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; @@ -68,42 +69,46 @@ bool ComposeTransposePermutations(const Graph&, Match& match) { ``` The pattern requires explicit `perm` attributes; it does not match default -reverse-axis Transposes. Known rank is enough: dimension sizes need not be +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. Optional attribute captures such as `p:ints | None` are not +supported: `None` describes tensor-slot/branch absence, not a missing attribute. +To handle an omitted attribute, leave it out of the pattern's keyword constraints +and inspect the matched node's `GetAttributes()` map in the checker; see +[attribute captures](#optional-type-and-shape-annotations) below for the default +attribute distinction. 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]`. -## 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. - ## 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 the first +successful concrete plan in source order, 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). +It does not run a checker or search another candidate after a later checker +rejection. Use the ordered visitor when that retry is needed: + ```cpp Pattern pattern("(x)=>{a=Identity(x);return Sigmoid(a)}"); size_t count = 0; @@ -125,10 +130,10 @@ 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. -The standalone API uses the same structural, attribute, constant, type/shape, -layer and **deletion-boundary** matcher as the engine. It enumerates rewrite-safe -patterns, not arbitrary subgraph occurrences. It deliberately does not apply -rule opset-range or execution-provider policy and does not run a checker, +Both standalone APIs use the same structural, attribute, constant, type/shape, +layer and **deletion-boundary** matcher as the engine. They find rewrite-safe +patterns, not arbitrary subgraph occurrences. They deliberately do not apply +rule opset-range or execution-provider policy and do not run a checker, construct a replacement or traverse other roots/subgraphs. The first consistent annotation combination is selected within each concrete graph plan; rejecting that candidate does not search other annotation combinations. @@ -282,6 +287,13 @@ 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, carry shapes, or use type unions (including `None`). Undeclared captures remain supported and untyped. +There is no nullable `optionalAttributeOf` lookup: `attributeOf` requires a bound +capture or prepared attribute and returns `const AttributeProto&`. A constrained +attribute must exist for a match; a schema default inserted by `Graph::Resolve` +counts as present. An optional attribute without a schema default (for example +Transpose `perm`) may remain absent. Omit its keyword constraint and use +`match.nodes[index]->GetAttributes().find("perm")` to distinguish that case in +C++; omitted/default permutations still need their operator-specific semantics. Annotations are constraints on **patterns only**; replacements continue to use unannotated inputs and derive types/shapes from preserved outputs and schema @@ -555,7 +567,7 @@ 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 `VisitMatches` and engine callbacks retain their selected compiled +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 @@ -565,7 +577,7 @@ 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 -`VisitMatches` or engine callbacks; author checkers using lookups, `update` and +`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. @@ -671,6 +683,29 @@ 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: @@ -894,7 +929,7 @@ normally and are not converted to nonmatches. `onnxruntime/test/optimizer/pattern_rewriter_components_test.cc` separates `PatternComponentTests`, `ReplacementComponentTests`, `MatchComponentTests` -and `CheckAndBindComponentTests`. These invoke `VisitMatches`, `Instantiate`, +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. diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 6248f5525fdbe..42b13340f90f2 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -1978,6 +1978,16 @@ Rule::Rule(std::string name, std::string_view pattern, std::string_view replacem 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."); diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index f61ca6cae1c6b..7ac908e5998fc 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -96,6 +96,8 @@ class Pattern { 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; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc index 990096e330c4e..968ad8f422ac7 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc @@ -148,6 +148,59 @@ TEST(PatternComponentTests, InvalidContextAndCallbackErrorsAreNotNonmatches) { 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, AttributeCapturesRequirePresenceAndDoNotSupportNone) { + EXPECT_THROW((void)Pattern("(x,p:ints|None)=>return Transpose(x,perm=p)"), OnnxRuntimeException); + 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); +} + TEST(PatternComponentTests, DeletionBoundaryButNotProviderPolicyIsIncluded) { for (int scenario = 0; scenario < 3; ++scenario) { TestGraph test; From ac7d3b1de01228efca9ca8b1e7b0dcacda3843a9 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Mon, 26 Oct 2026 10:43:34 -0700 Subject: [PATCH 41/45] Use readable whitespace throughout rewriter documentation examples Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 220 +++++++++++++++++++-------------------- 1 file changed, 110 insertions(+), 110 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 1a3a666b73390..7326c15bbce83 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -4,7 +4,7 @@ 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. +The following example replaces `Transpose(Transpose(x, p), q)` with one Transpose. ```cpp #include "core/optimizer/pattern_rewriter.h" @@ -19,11 +19,11 @@ bool ComposeTransposePermutations(const Graph&, Match& match) { Rule MakeComposeTransposeRule() { Pattern pattern(R"( - (x,p:ints,q:ints) => { - a=Transpose(x,perm=p) - return Transpose(a,perm=q) + (x, p : ints, q : ints) => { + a = Transpose(x, perm = p) + return Transpose(a, perm = q) })"); - Replacement replacement("(x,composed)=>Transpose(x,perm=composed)"); + Replacement replacement("(x, composed) => Transpose(x, perm = composed)"); return Rule{"ComposeTranspose", pattern, replacement, ComposeTransposePermutations, {{kOnnxDomain, 13, 25}}}; } @@ -39,7 +39,7 @@ 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]]`. +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: @@ -72,7 +72,7 @@ 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. Optional attribute captures such as `p:ints | None` are not +checker runs. Optional attribute captures such as `p : ints | None` are not supported: `None` describes tensor-slot/branch absence, not a missing attribute. To handle an omitted attribute, leave it out of the pattern's keyword constraints and inspect the matched node's `GetAttributes()` map in the checker; see @@ -82,7 +82,7 @@ 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]`. +noncommuting permutations `[1, 0, 2]` then `[0, 2, 1]` compose to `[1, 2, 0]`. ## Patterns @@ -91,7 +91,7 @@ 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)}"); +Pattern pattern("(x) => { a = Identity(x); return Sigmoid(a) }"); std::optional match; ORT_RETURN_IF_ERROR(pattern.MatchAt(graph, root.Index(), match)); if (match) { @@ -110,7 +110,7 @@ It does not run a checker or search another candidate after a later checker rejection. Use the ordered visitor when that retry is needed: ```cpp -Pattern pattern("(x)=>{a=Identity(x);return Sigmoid(a)}"); +Pattern pattern("(x) => { a = Identity(x); return Sigmoid(a) }"); size_t count = 0; ORT_RETURN_IF_ERROR(pattern.VisitMatches(graph, root_index, [&](Match& match) { ++count; @@ -182,8 +182,8 @@ 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)"); +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` @@ -202,13 +202,13 @@ are available for introspection. ### Optional type and shape annotations -Pattern header variables and assigned outputs may carry `: type` 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); + (x : float[M, N], y : float[N, K]) => { + product : float[M, K] = MatMul(x, y); return Identity(product); } )"); @@ -222,19 +222,19 @@ They cannot collide with tensor, attribute, or element-type variable names. |---|---| | `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` | +| `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. `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 +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. @@ -251,9 +251,9 @@ 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 +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. +`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 @@ -276,8 +276,8 @@ with computed attributes, and survive clearing computed attributes. 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)"); +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. @@ -306,8 +306,8 @@ silently ignored. Each annotation allows at most 64 dimensions, with at most 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)"); +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 @@ -322,7 +322,7 @@ The product of alternative counts across annotations is limited to 64 at parse t Empty branches, duplicate constraints, and `None` alone are rejected. `None | float[M]` is also accepted, but `None` is permitted only for header tensor inputs, not assigned outputs or attribute captures. Replacements remain untyped. -The former experimental `x?` syntax is removed; use `x:tensor[...] | None` for +The former experimental `x?` syntax is removed; use `x : tensor[...] | None` for an unconstrained optional tensor slot. ### Optional input slots @@ -332,13 +332,13 @@ 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); }"); +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 +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` @@ -358,11 +358,11 @@ 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])"); +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 +`[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: @@ -373,11 +373,11 @@ 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); }"); +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. +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 @@ -387,8 +387,8 @@ 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])"); +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`. @@ -402,8 +402,8 @@ 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)"); +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 @@ -427,10 +427,10 @@ 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) + (x, y) => { + sum = Add(x, y) + product = Mul(sum, x) + return Add(sum, product) } )"); ``` @@ -441,7 +441,7 @@ 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); }"); +Pattern pattern("(x) => { first, second = Split(x, axis = 0); return Concat(first, second, axis = 0); }"); Replacement replacement("(x) => Identity(x)"); ``` @@ -451,11 +451,11 @@ For a multi-output root, assign its outputs first, then return their names: Pattern pattern(R"( (x) => { copied = Identity(x); - left,right = Split(copied,axis=0); - return left,right; + left, right = Split(copied, axis = 0); + return left, right; } )"); -Replacement replacement("(x) => { left,right = Split(x,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, @@ -466,8 +466,8 @@ 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; }"); +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. @@ -498,7 +498,7 @@ The standalone `Instantiate` API executes a recipe without a pattern, Match, rule, or engine: ```cpp -Replacement recipe("(x,perm)=>Transpose(x,perm=perm)"); +Replacement recipe("(x, perm) => Transpose(x, perm = perm)"); ONNX_NAMESPACE::AttributeProto permutation; permutation.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INTS); permutation.add_ints(1); @@ -596,9 +596,9 @@ obtained from `VisitMatches`, without applying any rule: ```cpp std::optional accepted; ORT_RETURN_IF_ERROR(pattern.VisitMatches(graph, root_index, [&](Match& candidate) { - if (!ComposeTransposePermutations(graph, candidate)) return true; // Reject: try next. + if (!ComposeTransposePermutations(graph, candidate)) return true; // Reject : try next. accepted = candidate; // Keeps prepared "composed" plus selected-plan ownership. - return false; // Accept: stop enumeration. + return false; // Accept : stop enumeration. })); ``` @@ -612,9 +612,9 @@ 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 +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 resulting `FusedMatMul` uses `transB = 1` to perform the same transpose internally. The Microsoft-domain opset must already be imported. ```cpp @@ -625,9 +625,9 @@ using namespace onnxruntime::rewrite; Rule MakeTransposeMatMulRule() { const auto* pattern = R"( - (x:float[M,N], y:float[K,N]) => { + (x : float[M, N], y : float[K, N]) => { transposed = Transpose(y) - return MatMul(x,transposed) + return MatMul(x, transposed) } )"; auto check_and_bind = [](const Graph&, Match& match) { @@ -644,7 +644,7 @@ Rule MakeTransposeMatMulRule() { 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)"; + 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}}}; } ``` @@ -743,7 +743,7 @@ callers construct their own `Rules` and pass them to `PatternRewriter`. 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 +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 @@ -755,7 +755,7 @@ Contrib-disabled tests construct only the two HardSwish examples. ## Opset and execution-provider applicability -Each `OpsetRequirement` is `{domain, min_version, max_version}`, with positive, +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. @@ -838,8 +838,8 @@ 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 +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 @@ -903,7 +903,7 @@ 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'; + std::cerr << failure.rule << " : " << failure.reason << '\n'; }; ``` @@ -1003,7 +1003,7 @@ 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 +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. @@ -1019,25 +1019,25 @@ 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) + 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) + a = Mul(g, s) + product = Mul(a, u) return product } - pattern Fuse(g,u,alpha:float[...] | None) { - product = activation(g,u,alpha) + 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 +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 @@ -1049,21 +1049,21 @@ 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}`. +`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. +`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 +`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. @@ -1072,7 +1072,7 @@ 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 +`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. @@ -1082,18 +1082,18 @@ 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 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 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) +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,_) +pattern discovered(g, u) { + product = plain(g, u) | scaled(g, u, _) return product } ``` @@ -1127,9 +1127,9 @@ callee-local scope or deeper `alternativeView()`; grouped branches retain their nested selection metadata under the `` tag. 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 +`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.). @@ -1156,26 +1156,26 @@ may contain useful callable patterns that cannot independently serve as rules: ```text pattern identity(x) -> (y) { return x } -pattern branches(x,z) { +pattern branches(x, z) { left = Identity(x) right = Neg(z) - return left,right + return left, right } -pattern join(x,z) { - left,right = branches(x,z) +pattern join(x, z) { + left, right = branches(x, z) value = identity(left) - return Add(value,right) + 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)=>...")` +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 }` +`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 @@ -1190,8 +1190,8 @@ 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`. +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()`; @@ -1199,8 +1199,8 @@ 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 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 @@ -1221,7 +1221,7 @@ 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 +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. @@ -1377,9 +1377,9 @@ Rule MakePatternFusion(const PatternLibrary& library) { return true; }; Replacement recipe(R"( - (g,u,activation_alpha) => { - a = com.microsoft.QuickGelu(g,alpha=activation_alpha) - return Mul(a,u) + (g, u, activation_alpha) => { + a = com.microsoft.QuickGelu(g, alpha = activation_alpha) + return Mul(a, u) } )"); return Rule{"PatternActivation", library.GetPattern("Fuse"), recipe, check, @@ -1396,7 +1396,7 @@ 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)`. +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 @@ -1404,7 +1404,7 @@ when accessed through an interface alias. ### Optional immutable tensor constraints and constant readers -The prefix qualifier `const`, as in `top_k:const int64[1]`, requires an +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. @@ -1414,9 +1414,9 @@ It can qualify individual tensor alternatives, including `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 +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`. +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 From bfc52c7148ede6cf68915e96313c1f6d9592a0a8 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Mon, 26 Oct 2026 10:45:15 -0700 Subject: [PATCH 42/45] Keep diagnostic text and C++ punctuation unchanged in examples Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 7326c15bbce83..c987f09a97a3a 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -596,9 +596,9 @@ obtained from `VisitMatches`, without applying any rule: ```cpp std::optional accepted; ORT_RETURN_IF_ERROR(pattern.VisitMatches(graph, root_index, [&](Match& candidate) { - if (!ComposeTransposePermutations(graph, candidate)) return true; // Reject : try next. + if (!ComposeTransposePermutations(graph, candidate)) return true; // Reject: try next. accepted = candidate; // Keeps prepared "composed" plus selected-plan ownership. - return false; // Accept : stop enumeration. + return false; // Accept: stop enumeration. })); ``` @@ -755,7 +755,7 @@ Contrib-disabled tests construct only the two HardSwish examples. ## Opset and execution-provider applicability -Each `OpsetRequirement` is `{domain, min_version, max_version }`, with positive, +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. @@ -838,7 +838,7 @@ 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 +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. @@ -903,7 +903,7 @@ 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'; + std::cerr << failure.rule << ": " << failure.reason << '\n'; }; ``` From 6d707086e04fe287f162b9dc2baf58b935d8e2bc Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Mon, 26 Oct 2026 11:33:01 -0700 Subject: [PATCH 43/45] Explain ordered match visitation after pattern disjunction Keep the initial component API example focused on MatchAt, then introduce multiple candidates and checker retry with a concrete disjunction example. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 120 +++++++++++++++++++++++++-------------- 1 file changed, 78 insertions(+), 42 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index c987f09a97a3a..9d278ecd88cbb 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -101,42 +101,21 @@ if (match) { } ``` -`MatchAt` matches at exactly the specified root node. It returns the first -successful concrete plan in source order, or `std::nullopt` for an ordinary -nonmatch (clearing any previous result). Status errors and configuration +`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). -It does not run a checker or search another candidate after a later checker -rejection. Use the ordered visitor when that retry is needed: +The graph must already be resolved; an unresolved graph or invalid root returns +an error. Malformed/non-rooted pattern configuration throws. -```cpp -Pattern pattern("(x) => { a = Identity(x); return Sigmoid(a) }"); -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. -})); -``` - -`VisitMatches` anchors at one existing root in an **already resolved** graph. -It visits successful expanded graph plans in source order, each with fresh -bindings. `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. - -Both standalone APIs use the same structural, attribute, constant, type/shape, -layer and **deletion-boundary** matcher as the engine. They find rewrite-safe -patterns, not arbitrary subgraph occurrences. They deliberately do not apply -rule opset-range or execution-provider policy and do not run a checker, -construct a replacement or traverse other roots/subgraphs. The first consistent -annotation combination is selected within each concrete graph plan; rejecting -that candidate does not search other annotation combinations. +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 @@ -591,19 +570,21 @@ Declare the attribute binding in the replacement's input list, set it with 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 `VisitMatches`, without applying any rule: +obtained from `MatchAt`, without applying any rule: ```cpp -std::optional accepted; -ORT_RETURN_IF_ERROR(pattern.VisitMatches(graph, root_index, [&](Match& candidate) { - if (!ComposeTransposePermutations(graph, candidate)) return true; // Reject: try next. - accepted = candidate; // Keeps prepared "composed" plus selected-plan ownership. - return false; // Accept: stop enumeration. -})); +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; +} ``` -Do **not** directly return the checker boolean from a visitor: its traversal -meaning is opposite. In engine use, checker false discards candidate-local +For retry across disjunctions, see the later +[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 @@ -1277,6 +1258,61 @@ 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): + +```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. + ### Instance binding views and handwritten C++ structs `match.instance("product")` returns a read-only `BindingView` for the activation From e5927feb856711fc0f4c06145a553c1c2675cdf9 Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Mon, 26 Oct 2026 14:06:19 -0700 Subject: [PATCH 44/45] Support optional attribute captures and replacement omission Add typed attribute-or-None captures with explicit presence state, nullable scoped lookup, repeated-binding checks, and omission through shared replacement staging. Preserve standalone missing-name errors with explicit absent-attribute bindings. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 114 ++++++-- .../core/optimizer/pattern_rewriter.cc | 78 ++++-- onnxruntime/core/optimizer/pattern_rewriter.h | 10 +- .../pattern_rewriter_components_test.cc | 258 +++++++++++++++++- .../test/optimizer/pattern_rewriter_test.cc | 2 - 5 files changed, 418 insertions(+), 44 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 9d278ecd88cbb..857a1b7e38712 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -72,12 +72,11 @@ 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. Optional attribute captures such as `p : ints | None` are not -supported: `None` describes tensor-slot/branch absence, not a missing attribute. -To handle an omitted attribute, leave it out of the pattern's keyword constraints -and inspect the matched node's `GetAttributes()` map in the checker; see -[attribute captures](#optional-type-and-shape-annotations) below for the default -attribute distinction. Known rank is enough: dimension sizes need not be +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 @@ -212,7 +211,7 @@ They cannot collide with tensor, attribute, or element-type variable names. 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. -`None` in a type union describes an absent operator slot, not an ONNX optional +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. @@ -264,15 +263,10 @@ 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, carry -shapes, or use type unions (including `None`). Undeclared captures remain supported and untyped. -There is no nullable `optionalAttributeOf` lookup: `attributeOf` requires a bound -capture or prepared attribute and returns `const AttributeProto&`. A constrained -attribute must exist for a match; a schema default inserted by `Graph::Resolve` -counts as present. An optional attribute without a schema default (for example -Transpose `perm`) may remain absent. Omit its keyword constraint and use -`match.nodes[index]->GetAttributes().find("perm")` to distinguish that case in -C++; omitted/default permutations still need their operator-specific semantics. +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 @@ -299,8 +293,8 @@ 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, but `None` is permitted only for header tensor -inputs, not assigned outputs or attribute captures. Replacements remain untyped. +`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. @@ -393,8 +387,69 @@ 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`. Missing attributes -are not captured from schema defaults: only the resolved node's map is used. +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 @@ -497,6 +552,21 @@ 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 @@ -534,7 +604,7 @@ host C++ choices, not ONNX control flow. `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`, -`dimensionOf`, `elementTypeOf` and `BindingView` provide read-only access to captures. +`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 @@ -1317,7 +1387,7 @@ other annotation combinations. `match.instance("product")` returns a read-only `BindingView` for the activation call in the example. `view.instance("scaled")` selects its nested expression -occurrence. Its `valueOf`, `optionalValueOf`, `attributeOf`, `dimensionOf`, +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; diff --git a/onnxruntime/core/optimizer/pattern_rewriter.cc b/onnxruntime/core/optimizer/pattern_rewriter.cc index 42b13340f90f2..8189169c51916 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.cc +++ b/onnxruntime/core/optimizer/pattern_rewriter.cc @@ -66,7 +66,8 @@ 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)) { - return match.attributeOf(constraint.value_name).type() == constraint.attribute_type; + 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; @@ -199,7 +200,12 @@ bool MatchPattern(const Pattern& pattern, const Graph& graph, const Node& root, match.nodes[i] = node; for (const auto& attribute : expected.attributes) { auto found = node->GetAttributes().find(attribute.name); - if (found == node->GetAttributes().end() || !attribute.Matches(found->second)) + 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."); @@ -297,7 +303,7 @@ void RollbackFragment(Graph& graph, InlinedVector& added, const Inlin Status StageReplacement(Graph& graph, const Replacement& recipe, gsl::span returned_outputs, const std::function& resolve_tensor, - const std::function& resolve_attribute, + const std::function&)>& resolve_attribute, const Node* annotation_source, InlinedVector& added, InlinedVector& initializers) { ORT_RETURN_IF(!recipe.Plan(), "Replacement has no parsed plan."); @@ -320,10 +326,12 @@ Status StageReplacement(Graph& graph, const Replacement& recipe, gsl::span value; ORT_RETURN_IF_ERROR(resolve_attribute(binding, value)); - value.set_name(name); - attributes.emplace(name, std::move(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; @@ -454,8 +462,12 @@ Status Replace(Graph& graph, const Pattern& pattern, const Match& match, const R result = match.values[value] ? graph.GetNodeArg(match.values[value]->Name()) : nullptr; return Status::OK(); }; - auto attribute = [&](const std::string& name, ONNX_NAMESPACE::AttributeProto& result) -> Status { - result = match.attributeOf(name); + 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; @@ -1022,7 +1034,7 @@ class PatternParser { if (token_ != "|") break; ORT_RETURN_IF_ERROR(Next()); } while (true); - ORT_RETURN_IF(alternatives.size() == 1 && alternatives.front().absent, Error("None requires a tensor alternative")); + ORT_RETURN_IF(alternatives.size() == 1 && alternatives.front().absent, Error("None requires a concrete type alternative")); return Status::OK(); } @@ -1226,7 +1238,8 @@ class PatternParser { attribute_inputs_.insert(declared->second.handle); } pattern_.attribute_names_.insert(capture); - attributes.push_back({name, [](const auto&) { return true; }, 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_)); @@ -1392,8 +1405,12 @@ Status PatternParser::ValidateTypes(const Pattern& pattern) { combinations *= alternatives.size(); for (const auto& type : alternatives) { if (pattern.IsAttribute(handle)) { - ORT_RETURN_IF(alternatives.size() != 1 || !type.attribute_type || type.immutable || pattern.IsOptional(handle), - "Invalid attribute annotation or optional attribute."); + 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."); } @@ -2019,16 +2036,20 @@ Status Pattern::VisitMatches(const Graph& graph, NodeIndex root_index, } Status Replacement::Instantiate(Graph& graph, const TensorBindings& tensors, const NodeAttributes& attributes, - Instantiation& result, const Node* annotation_source) const { + 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), "Replacement binding is both tensor and attribute: ", name); + 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."); @@ -2037,7 +2058,11 @@ Status Replacement::Instantiate(Graph& graph, const TensorBindings& tensors, con value = found->second ? graph.GetNodeArg(found->second->Name()) : nullptr; return Status::OK(); }; - auto attribute = [&](const std::string& name, ONNX_NAMESPACE::AttributeProto& value) -> Status { + 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; @@ -2099,6 +2124,7 @@ Replacement::Replacement(std::string_view text) { 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); @@ -2123,6 +2149,10 @@ const NodeArg* BindingView::optionalValueOf(std::string_view name) const { 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)); @@ -2207,18 +2237,26 @@ void Match::update(std::string_view name, ONNX_NAMESPACE::AttributeProto 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 (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; + 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)); @@ -2227,6 +2265,12 @@ bool Match::BindAttribute(std::string_view name, const ONNX_NAMESPACE::Attribute 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); diff --git a/onnxruntime/core/optimizer/pattern_rewriter.h b/onnxruntime/core/optimizer/pattern_rewriter.h index 7ac908e5998fc..df8be925cdbe4 100644 --- a/onnxruntime/core/optimizer/pattern_rewriter.h +++ b/onnxruntime/core/optimizer/pattern_rewriter.h @@ -34,6 +34,7 @@ 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}); } @@ -161,6 +162,8 @@ class BindingView { 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; @@ -194,6 +197,7 @@ struct Match { 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, ""); } @@ -203,6 +207,7 @@ struct Match { 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) { @@ -217,6 +222,7 @@ struct Match { 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_; @@ -252,8 +258,10 @@ struct Replacement { 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; + Instantiation& result, const Node* annotation_source = nullptr, + const InlinedHashSet& absent_attributes = {}) const; private: friend class PatternRewriter; diff --git a/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc index 968ad8f422ac7..6cb78e86dc988 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_components_test.cc @@ -184,8 +184,7 @@ TEST(PatternComponentTests, MatchAtReturnsFirstCandidateAndClearsOrdinaryNonmatc ASSERT_TRUE(result); } -TEST(PatternComponentTests, AttributeCapturesRequirePresenceAndDoNotSupportNone) { - EXPECT_THROW((void)Pattern("(x,p:ints|None)=>return Transpose(x,perm=p)"), OnnxRuntimeException); +TEST(PatternComponentTests, OptionalAttributeCapturesDistinguishAbsenceFromUnknownNames) { TestGraph test; auto& b = test.builder; auto* x = b.MakeInput({2, 3}, -1.0f, 1.0f); @@ -199,6 +198,261 @@ TEST(PatternComponentTests, AttributeCapturesRequirePresenceAndDoNotSupportNone) 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) { diff --git a/onnxruntime/test/optimizer/pattern_rewriter_test.cc b/onnxruntime/test/optimizer/pattern_rewriter_test.cc index e771747f79ee7..c363d2809c64d 100644 --- a/onnxruntime/test/optimizer/pattern_rewriter_test.cc +++ b/onnxruntime/test/optimizer/pattern_rewriter_test.cc @@ -2906,7 +2906,6 @@ TEST(PatternRewriterTests, TypeUnionParsingAndBounds) { "(x?)=>return Identity(x)", "(x?:float[])=>return Identity(x)", "(x:float[]?)=>return Identity(x)", "(x,p:ints|floats)=>return Identity(x,perm=p)", - "(x,p:int|None)=>return Identity(x,axis=p)", "(x)=>{y:float[]|None=Identity(x);return y;}"}) { SCOPED_TRACE(text); EXPECT_THROW((void)Pattern{text}, OnnxRuntimeException); @@ -6118,7 +6117,6 @@ TEST(PatternRewriterTests, OrdinaryOutputSignaturesPreserveTypedAttributeInputs) 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|None)->(y){y=Softmax(x,axis=a);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); From 319fa38a679b3730ca7487c2269fa388dca58f2b Mon Sep 17 00:00:00 2001 From: Ganesan Ramalingam Date: Mon, 26 Oct 2026 15:05:31 -0700 Subject: [PATCH 45/45] Present pattern libraries immediately after patterns Keep the composed checker/replacement example as a later standalone section and update navigation for the new reading order. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: daf231a0-ce61-4a2b-9649-0f31e2296e4f --- docs/Pattern_Rewriter.md | 1100 +++++++++++++++++++------------------- 1 file changed, 553 insertions(+), 547 deletions(-) diff --git a/docs/Pattern_Rewriter.md b/docs/Pattern_Rewriter.md index 857a1b7e38712..95f21704e9fb8 100644 --- a/docs/Pattern_Rewriter.md +++ b/docs/Pattern_Rewriter.md @@ -520,6 +520,361 @@ 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 @@ -652,7 +1007,7 @@ if (candidate && ComposeTransposePermutations(graph, *candidate)) { } ``` -For retry across disjunctions, see the later +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 @@ -823,570 +1178,246 @@ nodes are eligible; assigned nodes must belong to 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. +```cpp +Options options; +options.compatible_execution_providers.insert(kCudaExecutionProvider); +PatternRewriter rewriter(Rules{MakeComposeTransposeRule()}, options); +``` -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. +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. -## Composable pattern libraries +## The rewriter engine: applying collections of rules -`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 every -existing `Rule` constructor. The library and its source text need not outlive -the returned pattern. +Register the transformer through the normal session pipeline: ```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 - } +auto transformer = std::make_unique(Rules{MakeComposeTransposeRule()}); +ORT_RETURN_IF_ERROR(session.RegisterGraphTransformer( + std::move(transformer), TransformerLevel::Level1)); +``` - pattern Fuse(g, u, alpha : float[...] | None) { - product = activation(g, u, alpha) - return Identity(product) - } -)"); +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()); ``` -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. +`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. -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. +### Resolved metadata and mutation boundaries -`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. +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. -### Expressions, alternatives, and branch-local inputs +`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. -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. +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. -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. +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. -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: +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. -```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 -} -``` +`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. -`_` 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. +```cpp +Options options; +options.max_iterations = 3; +PatternRewriter rewriter(rules, options); +``` -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. +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. -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. +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. -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.). +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. -### Calls versus standalone entries +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. -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: +`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. -```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) -} +### 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'; +}; ``` -`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. +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. -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. +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. -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: +## 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 -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) -} +onnxruntime_test_all --gtest_filter=PatternComponentTests.*:ReplacementComponentTests.*:MatchComponentTests.*:CheckAndBindComponentTests.*:PatternRewriterTests.* ``` -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. +The component split makes contracts independently testable; it does not bypass +the engine's integration tests or duplicate a second matcher/recipe executor. -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. +## Limits and unsupported forms -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. +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`. -Each complete candidate passes the existing structural, type/constant, -deletion-boundary, layer, provider and opset checks 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. +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. -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 algorithm outline -### Matching alternatives at one root +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. -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. +Multiple rules are handled **node-first**, inside an outer bounded retry loop: -```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. -})); +```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 ``` -`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): - -```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. -})); -``` +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. -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. +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 +## Instance binding views and handwritten C++ structs `match.instance("product")` returns a read-only `BindingView` for the activation -call in the example. `view.instance("scaled")` selects its nested expression +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 @@ -1507,28 +1538,3 @@ 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. - -### 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.