[Rewriter] Adopt missing ONNX Light optimization patterns by justinchuby · Pull Request #3167 · microsoft/onnxscript

Justin Chu (justinchuby) · 2026-11-05T23:42:37Z

Add or extend sound coverage for 69 upstream registrations across canonicalization, collections, Boolean selection, shapes/layouts, expansion, matrix/attention, and LabelEncoder rules; wire them into the default optimizer.
Account for all 129 registrations at 43c2ac6, with pinned sources, licensing, existing equivalents, exact subsets, and exclusion reasons in docs/articles/onnx_light_patterns.md.
Preserve existing ONNX/ORT rotary and attention fusion opportunities after canonicalization, with regression coverage for broadcasting, dtypes/opsets, shared values, overridable initializers, and numerical edge cases.

Validation: full onnxscript/rewriter and onnxscript/optimizer suites: 867 passed, 3 existing skips, 107 subtests passed. lintrunner -a and git diff --check clean.

Inventory all 129 registered patterns at upstream revision 43c2ac6aa6a393dfb7b90527638d0c1f22139c1e and independently implement sound missing subsets. Preserve existing rotary and attention fusion opportunities after canonicalization. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

+    def pattern(self, op, x0, x1, axis):
+        return op.Concat(x0, x1, axis=axis)
+
+    def check(self, context, x0, x1, axis):


+    def check(self, context, x0, x1, axis):
+        return super().check(context, axis, x0=x0, x1=x1)
+
+    def rewrite(self, op, x0, x1, axis):


+    def pattern(self, op, x0, x1, x2, axis):
+        return op.Concat(x0, x1, x2, axis=axis)
+
+    def check(self, context, x0, x1, x2, axis):


+    def check(self, context, x0, x1, x2, axis):
+        return super().check(context, axis, x0=x0, x1=x1, x2=x2)
+
+    def rewrite(self, op, x0, x1, x2, axis):


+    def pattern(self, op, x0, x1, x2, x3, axis):
+        return op.Concat(x0, x1, x2, x3, axis=axis)
+
+    def check(self, context, x0, x1, x2, x3, axis):


+    def check(
+        self,
+        context,
+        x,
+        s0,
+        e0,
+        a0,
+        st0,
+        s1,
+        e1,
+        a1,
+        st1,
+        s2,
+        e2,
+        a2,
+        st2,
+    ):


+    def pattern(
+        self,
+        op,
+        x,
+        s0,
+        e0,
+        a0,
+        st0,
+        s1,
+        e1,
+        a1,
+        st1,
+        s2,
+        e2,
+        a2,
+        st2,
+        s3,
+        e3,
+        a3,
+        st3,
+    ):


+    def check(
+        self,
+        context,
+        x,
+        s0,
+        e0,
+        a0,
+        st0,
+        s1,
+        e1,
+        a1,
+        st1,
+        s2,
+        e2,
+        a2,
+        st2,
+        s3,
+        e3,
+        a3,
+        st3,
+    ):


+    cast = next(node for node in optimized.graph if node.op_type == "Cast")
+    assert cast.attributes["to"].as_int() == 10
+    assert cast.attributes["saturate"].as_int() == 0
+    del original


+    assert op_types.index("Slice") < op_types.index("Cast")
+    cast = next(node for node in optimized.graph if node.op_type == "Cast")
+    assert cast.attributes["saturate"].as_int() == 0
+    del original


+    def pattern(self, op, x0, x1, axis):
+        return op.Concat(x0, x1, axis=axis)
+
+    def check(self, context, x0, x1, axis):


+    def check(self, context, x0, x1, axis):
+        return super().check(context, axis, x0=x0, x1=x1)
+
+    def rewrite(self, op, x0, x1, axis):


+    def pattern(self, op, x0, x1, x2, axis):
+        return op.Concat(x0, x1, x2, axis=axis)
+
+    def check(self, context, x0, x1, x2, axis):


+    def check(self, context, x0, x1, x2, axis):
+        return super().check(context, axis, x0=x0, x1=x1, x2=x2)
+
+    def rewrite(self, op, x0, x1, x2, axis):


+    def pattern(self, op, x0, x1, x2, x3, axis):
+        return op.Concat(x0, x1, x2, x3, axis=axis)
+
+    def check(self, context, x0, x1, x2, x3, axis):


+    def check(self, context, x0, x1, x2, x3, axis):
+        return super().check(context, axis, x0=x0, x1=x1, x2=x2, x3=x3)
+
+    def rewrite(self, op, x0, x1, x2, x3, axis):


+
+
+class GathersSplit3(GathersSplit2):
+    def pattern(self, op, x, i0, i1, i2, axis0, axis1, axis2):


+            op.Gather(x, i2, axis=axis2, _check=_gather_index_checker(2)),
+        )
+
+    def check(self, context, x, i0, i1, i2, axis0, axis1, axis2):


+
+
+class GathersSplit4(GathersSplit2):
+    def pattern(self, op, x, i0, i1, i2, i3, axis0, axis1, axis2, axis3):


+            op.Gather(x, i3, axis=axis3, _check=_gather_index_checker(3)),
+        )
+
+    def check(self, context, x, i0, i1, i2, i3, axis0, axis1, axis2, axis3):


codecov · 2026-11-05T23:50:41Z

Codecov Report

❌ Patch coverage is 80.13873% with 859 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.19%. Comparing base (5b210da) to head (15dec59).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
...rewriter/rules/common/_collection_optimizations.py	72.07%	166 Missing and 82 partials ⚠️
...ipt/rewriter/rules/common/_layout_optimizations.py	70.78%	115 Missing and 86 partials ⚠️
...ipt/rewriter/rules/common/_matrix_optimizations.py	74.20%	80 Missing and 42 partials ⚠️
...ipt/rewriter/rules/common/_expand_optimizations.py	76.43%	63 Missing and 23 partials ⚠️
.../rewriter/rules/common/_canonical_optimizations.py	77.81%	35 Missing and 36 partials ⚠️
...ript/rewriter/rules/common/_shape_optimizations.py	82.16%	16 Missing and 17 partials ⚠️
onnxscript/rewriter/ort_fusions/mha.py	60.29%	11 Missing and 16 partials ⚠️
...pt/rewriter/rules/common/_boolean_optimizations.py	83.89%	9 Missing and 10 partials ⚠️
...ewriter/rules/common/_matrix_optimizations_test.py	93.63%	16 Missing and 1 partial ⚠️
onnxscript/rewriter/rules/common/_label_encoder.py	79.59%	5 Missing and 5 partials ⚠️
... and 7 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3167      +/-   ##
==========================================
+ Coverage   73.34%   74.19%   +0.85%     
==========================================
  Files         269      283      +14     
  Lines       33564    37825    +4261     
  Branches     3210     3835     +625     
==========================================
+ Hits        24618    28066    +3448     
- Misses       7876     8381     +505     
- Partials     1070     1378     +308

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

github-project-automation Bot added this to ONNX Script Review Board Nov 5, 2026

github-project-automation Bot moved this to Todo in ONNX Script Review Board Nov 5, 2026

github-advanced-security Bot found potential problems Nov 5, 2026

View reviewed changes

github-code-quality Bot found potential problems Nov 5, 2026

View reviewed changes

Navigation Menu

Repository navigation

[Rewriter] Adopt missing ONNX Light optimization patterns - #3167

[Rewriter] Adopt missing ONNX Light optimization patterns#3167
Justin Chu (justinchuby) wants to merge 1 commit into
mainfrom
justinchuby-onnx-light-pattern-optimizations

Justin Chu (justinchuby) commented Nov 5, 2026

Uh oh!

codecov Bot commented Nov 5, 2026 •

edited

Loading

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		class GathersSplit3(GathersSplit2):
		def pattern(self, op, x, i0, i1, i2, axis0, axis1, axis2):



		class GathersSplit4(GathersSplit2):
		def pattern(self, op, x, i0, i1, i2, i3, axis0, axis1, axis2, axis3):

Conversation

Justin Chu (justinchuby) commented Nov 5, 2026

Uh oh!

codecov Bot commented Nov 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Nov 5, 2026 •

edited

Loading