[rewriter] Decouple llama rule sets and make API explicit #2388

Copilot · 2025-06-14T15:32:54Z

This PR addresses the misleading naming and tangled organization of rewrite rules by decoupling the llama_rule_sets.py module and creating a more explicit API.

Problem

The original llama_rule_sets.py contained general optimization rules that weren't specific to Llama models, making the naming misleading. The API didn't explicitly specify what rules were being applied, making it unclear what optimizations were happening.

# Before: Unclear what this does
from onnxscript.rewriter import llama_rule_sets
rules = llama_rule_sets.llama_p0_rule_set()  # What rules? Why "llama"? What's "p0"?

Solution

1. Created `basic_rules.py` with explicit naming

Moved all general optimization rules to a new basic_rules.py module
Used descriptive function name: basic_optimization_rules()
Added comprehensive documentation for each rule

2. Made API explicit for fine-grained control

# New explicit API - users know exactly what they're getting
from onnxscript.rewriter import basic_rules

# Use all basic optimizations (recommended default)
rules = basic_rules.basic_optimization_rules()

# Or use specific individual rules
transpose_rule = basic_rules.transpose_identity_rule
cast_rule = basic_rules.cast_identity_rule

# Or create custom rule combinations
custom_rules = basic_rules.orp.RewriteRuleSet([
    basic_rules.transpose_identity_rule,
    basic_rules.cast_identity_rule,
])

3. Updated default rewriter to be explicit

# Before (in rewriter/__init__.py)
*llama_rule_sets.llama_p0_rule_set().rules,

# After - much clearer what's being applied
*basic_rules.basic_optimization_rules().rules,

4. Maintained backward compatibility

llama_rule_sets.py now serves as a compatibility wrapper
All existing APIs continue to work with deprecation warnings
Existing tests pass unchanged

Available Rules

The new API provides access to these optimization rules:

cast_cast_rule - Eliminates consecutive casts
cast_identity_rule - Removes redundant casts
expand_identity_rule - Removes no-op expands
reshape_reshape_rule - Combines consecutive reshapes
slice_split_rule - Converts slices to splits when beneficial
transpose_identity_rule - Removes identity transposes
transpose_transpose_rule - Combines consecutive transposes
unsqueeze_unsqueeze_rule - Combines consecutive unsqueezes
squeeze_reshape_1d_rule - Optimizes 1D squeeze+reshape patterns

Migration

# OLD (deprecated but still works)
from onnxscript.rewriter import llama_rule_sets
rules = llama_rule_sets.llama_p0_rule_set()

# NEW (recommended)
from onnxscript.rewriter import basic_rules
rules = basic_rules.basic_optimization_rules()

This change resolves the core issue by making the optimizer API explicitly specify what rules are being applied, while providing users with fine-grained control over optimization behavior.

Fixes #2128.

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.

Co-authored-by: justinchuby <[email protected]>

justinchuby

Please rename onnxscript/rewriter/llama_rule_sets_test.py and remove onnxscript/rewriter/llama_rule_sets.py entirely

Co-authored-by: justinchuby <[email protected]>

Copilot · 2025-06-14T16:07:04Z

Please rename onnxscript/rewriter/llama_rule_sets_test.py and remove onnxscript/rewriter/llama_rule_sets.py entirely

Done! Removed onnxscript/rewriter/llama_rule_sets.py entirely and renamed onnxscript/rewriter/llama_rule_sets_test.py to deprecated_llama_rule_sets_test.py. The backward compatibility layer has been completely eliminated as requested. See commit 362095c.

justinchuby · 2025-06-14T16:10:28Z

onnxscript/rewriter/deprecated_llama_rule_sets_test.py

Just move everything to onnxscript/rewriter/basic_rules_test.py and remove this file

Done! Removed deprecated_llama_rule_sets_test.py entirely since all its test functionality was already present in basic_rules_test.py with better naming. All tests continue to pass. See commit af80d0d.

onnxscript/rewriter/basic_rules_test.py

codecov · 2025-06-14T16:13:19Z

❌ 3 Tests Failed:

Tests completed	Failed	Passed	Skipped
16441	3	16438	2361

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0250_test_clip_outbounds

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_clip_outbounds'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_clip_outbounds' (e=No module named 'tests.onnx_backend_test_code.test_clip_outbounds') (file: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_clip_outbounds.py', absolute path: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_clip_outbounds.py', current folder: C:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_clip_outbounds(x: FLOAT[3], min: FLOAT, max: FLOAT) -> (FLOAT[3]):
E       y = opset13.Clip(x, min, max)
E       return y

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0668_test_mod_mixed_sign_int16

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_mod_mixed_sign_int16'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_mod_mixed_sign_int16' (e=No module named 'tests.onnx_backend_test_code.test_mod_mixed_sign_int16') (file: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_mod_mixed_sign_int16.py', absolute path: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_mod_mixed_sign_int16.py', current folder: C:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import INT16
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_mod_mixed_sign_int16(x: INT16[6], y: INT16[6]) -> (INT16[6]):
E       z = opset13.Mod(x, y)
E       return z

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1104_test_shape_start_1_end_2

Stack Traces | 0.004s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
C:\hostedtoolcache\windows\Python\3.12.10\x64\Lib\importlib\__init__.py:90: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_shape_start_1_end_2'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_shape_start_1_end_2' (e=No module named 'tests.onnx_backend_test_code.test_shape_start_1_end_2') (file: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_shape_start_1_end_2.py', absolute path: 'C:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_shape_start_1_end_2.py', current folder: C:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset21
E   
E   @script()
E   def bck_test_shape_start_1_end_2(x: FLOAT[3,4,5]) -> (INT64[1]):
E       y = opset21.Shape(x, end=2, start=1)
E       return y

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

Co-authored-by: justinchuby <[email protected]>

onnxscript/rewriter/basic_rules.py

onnxscript/rewriter/basic_rules_test.py

titaiwangms

lint

Signed-off-by: Justin Chu <[email protected]>

onnxscript/rewriter/basic_rules_test.py

-        rule_set.apply_to_model(model)
-        rewritten_model = ir.serde.serialize_model(model)
+        rule.apply_to_model(model)
+        _rewritten_model = ir.serde.serialize_model(model)


To fix the issue, the assignment to _rewritten_model on line 160 should be removed entirely, as the variable is not used anywhere in the code. This ensures that the code is clean and avoids unnecessary assignments. Since the right-hand side of the assignment (ir.serde.serialize_model(model)) does not have any side effects, it can be safely removed without impacting the functionality of the code.

Initial plan for issue

74bf271

Copilot AI assigned Copilot and justinchuby Jun 14, 2025

github-project-automation bot added this to ONNX Script Review Board Jun 14, 2025

github-project-automation bot moved this to Todo in ONNX Script Review Board Jun 14, 2025

Copilot started work on behalf of justinchuby June 14, 2025 15:32 View session

Copilot AI and others added 2 commits June 14, 2025 15:40

Create basic_rules.py and update rewriter to use it

7ccdbfb

Co-authored-by: justinchuby <[email protected]>

Update llama_rule_sets_test.py to use basic_rules

f53863a

Co-authored-by: justinchuby <[email protected]>

Copilot AI changed the title ~~[WIP] [rewriter] Decouple llama rule sets~~ [rewriter] Decouple llama rule sets and make API explicit Jun 14, 2025

Copilot AI requested a review from justinchuby June 14, 2025 15:50

Copilot finished work on behalf of justinchuby June 14, 2025 15:50

justinchuby reviewed Jun 14, 2025

View reviewed changes

Copilot started work on behalf of justinchuby June 14, 2025 16:02 View session

Remove llama_rule_sets.py and rename test file as requested

362095c

Co-authored-by: justinchuby <[email protected]>

Copilot finished work on behalf of justinchuby June 14, 2025 16:07

justinchuby added the topic: breaking changes label Jun 14, 2025

justinchuby requested a review from titaiwangms June 14, 2025 16:09

justinchuby marked this pull request as ready for review June 14, 2025 16:09

justinchuby reviewed Jun 14, 2025

View reviewed changes

Copilot started work on behalf of justinchuby June 14, 2025 16:10 View session

github-advanced-security bot found potential problems Jun 14, 2025

View reviewed changes

onnxscript/rewriter/basic_rules_test.py Fixed Show fixed Hide fixed

Remove deprecated_llama_rule_sets_test.py as requested

af80d0d

Co-authored-by: justinchuby <[email protected]>

Copilot finished work on behalf of justinchuby June 14, 2025 16:15

github-advanced-security bot found potential problems Jun 16, 2025

View reviewed changes

titaiwangms approved these changes Jun 16, 2025

View reviewed changes

justinchuby added 3 commits June 17, 2025 09:40

Update

b9a0e04

Signed-off-by: Justin Chu <[email protected]>

test

4169da5

Signed-off-by: Justin Chu <[email protected]>

Merge branch 'main' into copilot/fix-2128

9388dd0

github-advanced-security bot found potential problems Jun 17, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[rewriter] Decouple llama rule sets and make API explicit #2388

[rewriter] Decouple llama rule sets and make API explicit #2388

Copilot AI commented Jun 14, 2025 •

edited

Loading

Uh oh!

justinchuby left a comment

Uh oh!

Copilot AI commented Jun 14, 2025

Uh oh!

justinchuby Jun 14, 2025

Uh oh!

Copilot AI Jun 14, 2025

Uh oh!

Uh oh!

codecov bot commented Jun 14, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

titaiwangms left a comment

Uh oh!

Check notice

Copilot Autofix

Uh oh!

[rewriter] Decouple llama rule sets and make API explicit #2388

Are you sure you want to change the base?

[rewriter] Decouple llama rule sets and make API explicit #2388

Conversation

Copilot AI commented Jun 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Solution

1. Created basic_rules.py with explicit naming

2. Made API explicit for fine-grained control

3. Updated default rewriter to be explicit

4. Maintained backward compatibility

Available Rules

Migration

Uh oh!

justinchuby left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI commented Jun 14, 2025

Uh oh!

justinchuby Jun 14, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jun 14, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

codecov bot commented Jun 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

❌ 3 Tests Failed:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

titaiwangms left a comment

Choose a reason for hiding this comment

Uh oh!

Check notice

Copilot Autofix

Uh oh!

Copilot AI commented Jun 14, 2025 •

edited

Loading

1. Created `basic_rules.py` with explicit naming

codecov bot commented Jun 14, 2025 •

edited

Loading