Pattern optimization#

onnx-light rewrites graphs with a pattern-based optimizer built directly on top of GraphBuilder. The optimizer recognizes local subgraphs and replaces them with cheaper equivalents, in the spirit of the Python pattern optimizer it is ported from. The implementation plan and the pull requests that delivered it are recorded in Pattern-based optimization in GraphBuilder.

Overview#

Optimization always operates on a GraphBuilder through GraphGraph. GraphGraph wraps a builder with a structural index (successors, predecessors, shape, type and constant queries) and drives a match/apply loop:

from onnx_light.onnx_core.optimization import GraphBuilder, GraphGraph

builder = GraphBuilder(model)
graph = GraphGraph(builder)
rewrites, report = graph.optimize(report=True)
optimized_model = builder.to_onnx("model")

Because the optimizer reuses the builder, it inherits the builder’s shape and type inference, its constant knowledge and its cleanup passes (RemoveIdentityNodes, RemoveUnusedNodes, RemoveDuplicateNodes) instead of duplicating them.

The rewrite invariant#

A pattern must never reuse an existing name: every value it produces is new. This invariant keeps the successor and predecessor maps valid between two rewrites of the same iteration, which is why the builder records every name it hands out and never reuses one.

Pattern registration#

Patterns use the same global-plus-local model as shape functions. Registries are merged by the stable PatternOptimization.name; a more local entry replaces an entry with the same name:

  • global patterns (register_pattern()) are used by every new GraphGraph; the standard ONNX patterns are registered globally when the module is imported;

  • builder patterns (GraphBuilder.register_pattern) override a global pattern for optimizers built over that builder;

  • graph patterns (GraphGraph(builder, patterns=[...])) have the highest precedence and are retained for that optimizer, including recursive subgraphs.

Patterns can be written in C++ or in Python; both share the PatternOptimization interface, a match step that returns a MatchResult and an apply step that produces the replacement nodes.

API reference#

Examples#