Skip to content

[IR] Improve pass infra #2120

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 21 commits into from
Mar 26, 2025
Merged

[IR] Improve pass infra #2120

merged 21 commits into from
Mar 26, 2025

Conversation

justinchuby
Copy link
Collaborator

@justinchuby justinchuby commented Mar 22, 2025

  1. Run invariant functions requires and ensures by default at Pass __call__ to match pytorch's pass behavior. This means the invariants cannot be too expensive because they are always checked.
  2. Make PassManager a Pass so that it can be composed.
  3. Add changes_input attribute to indicate if the input is changed. Turn two class attributes into properties for them to be dynamic. Combining the two attributes we can tell if a pass is destructive. For now the properties are unused but they will become useful when we want to have a better guard on pass usage etc.
  4. Create Sequential, InPlacePass, FunctionalPass to help users create passes.

Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR updates the pass infrastructure to add built‐in precondition and postcondition checks in PassBase and refactors the PassManager to inherit from PassBase.

  • Added a new attribute “destructive” to indicate whether the pass can destroy the input model.
  • Implemented try/except logic for precondition and postcondition checks in PassBase.call.
  • Refactored PassManager to be a subclass of PassBase and removed redundant invariant checking.

Copy link

codecov bot commented Mar 22, 2025

❌ 4 Tests Failed:

Tests completed Failed Passed Skipped
13898 4 13894 1901
View the top 3 failed test(s) by shortest run time
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0377_test_flatten_default_axis
Stack Traces | 0.005s run time
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.10.11\x64\lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_flatten_default_axis'

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_flatten_default_axis' (e=No module named 'tests.onnx_backend_test_code.test_flatten_default_axis') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_flatten_default_axis.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_flatten_default_axis.py', current folder: D:\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 opset21
E   
E   @script()
E   def bck_test_flatten_default_axis(a: FLOAT[5,4,3,2]) -> (FLOAT[5,24]):
E       b = opset21.Flatten(a)
E       return b
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0125_test_ai_onnx_ml_tree_ensemble_set_membership
Stack Traces | 0.012s run time
onnxscript/converter.py:467: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
E   NameError: name 'nan' is not defined

The above exception was the direct cause of the following exception:
..../test_ort_nightly/lib/python3.11.../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:137: in extract_functions
    mod = importlib.import_module(import_name)
.../Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
..../test_ort_nightly/lib/python3.11.../_pytest/assertion/rewrite.py:185: in exec_module
    exec(co, module.__dict__)
tests/onnx_backend_test_code/test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
onnxscript/main.py:91: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:35: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1460: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1447: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:969: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1056: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:1000: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:553: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:832: in _translate_call_expr
    attrs = [
onnxscript/converter.py:833: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:517: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:469: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0397_test_ai_onnx_ml_tree_ensemble_set_membership
Stack Traces | 0.024s run time
onnxscript/converter.py:467: in _eval_constant_expr
    return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
E   NameError: name 'nan' is not defined

The above exception was the direct cause of the following exception:
..../test_ort_nightly/lib/python3.11.../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:137: in extract_functions
    mod = importlib.import_module(import_name)
.../hostedtoolcache/Python/3.11.11.../x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
<frozen importlib._bootstrap>:1204: in _gcd_import
    ???
<frozen importlib._bootstrap>:1176: in _find_and_load
    ???
<frozen importlib._bootstrap>:1147: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:690: in _load_unlocked
    ???
..../test_ort_nightly/lib/python3.11.../_pytest/assertion/rewrite.py:185: in exec_module
    exec(co, module.__dict__)
tests/onnx_backend_test_code/test_ai_onnx_ml_tree_ensemble_set_membership.py:9: in <module>
    @script()
onnxscript/main.py:91: in transform
    result = script_check(f_ast, opset, env, src, default_opset=default_opset)
onnxscript/main.py:35: in script_check
    return convert.translate_function_def(f)
onnxscript/converter.py:1460: in translate_function_def
    fn_ir = self._translate_function_def_common(stmt)
onnxscript/converter.py:1447: in _translate_function_def_common
    self._translate_stmt(s, index_of_stmt=i)
onnxscript/converter.py:969: in _translate_stmt
    return self._translate_assign_stmt(node)
onnxscript/converter.py:1056: in _translate_assign_stmt
    assign(lhs, rhs)
onnxscript/converter.py:1000: in assign
    t = self._translate_expr(rhs, lhs).name
onnxscript/converter.py:553: in _translate_expr
    r = self._translate_call_expr(node)
onnxscript/converter.py:832: in _translate_call_expr
    attrs = [
onnxscript/converter.py:833: in <listcomp>
    self._translate_attr(x, y, callee.op_schema.attributes[x])
onnxscript/converter.py:517: in _translate_attr
    val = self._eval_constant_expr(expr)
onnxscript/converter.py:469: in _eval_constant_expr
    raise NameError(
E   NameError: ERROR: Missing names, globals contains ['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__file__', '__cached__', '__builtins__', '@py_builtins', '@pytest_ar', 'numpy', 'TensorProto', 'make_tensor', 'script', 'external_tensor', 'Opset', 'FLOAT', 'ai_onnx_ml5'], locals [].
E   at: Function 'bck_test_ai_onnx_ml_tree_ensemble_set_membership', line 3
E       Y = ai_onnx_ml5.TreeEnsemble(X, aggregate_function=1, leaf_targetids=[0, 1, 2, 3], leaf_weights=make_tensor("value", 1, dims=[4], vals=[1.0, 10.0, 1000.0, 100.0]), membership_values=make_tensor("value", 1, dims=[8], vals=[1.2000000476837158, 3.700000047683716, 8.0, 9.0, nan, 12.0, 7.0, nan]), n_targets=4, nodes_falseleafs=[1, 0, 1], nodes_falsenodeids=[2, 2, 3], nodes_featureids=[0, 0, 0], nodes_modes=make_tensor("value", 2, dims=[3], vals=[0, 6, 6]), nodes_splits=make_tensor("value", 1, dims=[3], vals=[11.0, 232344.0, nan]), nodes_trueleafs=[0, 1, 1], nodes_truenodeids=[1, 0, 1], post_transform=0, tree_roots=[0])
E                                                                                                                                                                                             ^

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

@justinchuby justinchuby added module: IR Intermediate representation topic: passes labels Mar 22, 2025
@justinchuby justinchuby changed the title [IR] Update pass infra [IR] Improve pass infra Mar 22, 2025
justinchuby and others added 3 commits March 21, 2025 19:52
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
@justinchuby justinchuby requested a review from Copilot March 22, 2025 02:55
Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This pull request updates the pass infrastructure to automatically run invariant functions (requires and ensures) on every pass call, makes PassManager a subclass of PassBase, and converts class attributes to dynamic properties.

  • Invariant checks (requires and ensures) are now executed by default in PassBase.call.
  • PassManager is refactored to subclass PassBase and remove the redundant check_invariants flag.
  • The attributes in_place and destructive were converted from static class attributes to dynamic properties.

@justinchuby
Copy link
Collaborator Author

@gramalingam

@justinchuby justinchuby requested a review from Copilot March 22, 2025 03:01
Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR enhances the pass infrastructure by enforcing invariant checks at the PassBase level and making the PassManager composable as a Pass. Key changes include:

  • Integrating default invariant checks (via requires and ensures) into PassBase.call.
  • Converting PassManager into a subclass of PassBase with dynamic in_place and destructive properties.
  • Removing the optional invariant flag and cleaning up redundant invariant checks.
Comments suppressed due to low confidence (1)

onnxscript/ir/passes/_pass_infra.py:105

  • The concatenated string for the TypeError message is missing a space between the two sentences, which could reduce clarity. Consider inserting a space or apostrophe between the string literals.
raise TypeError(
                f"The result of the pass '{self.__class__.__name__}' should be type PassResult.""Please create one with ir.passes.PassResult()."

@justinchuby justinchuby requested a review from Copilot March 22, 2025 03:22
Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR improves the pass infrastructure by enforcing invariant checks systematically, making the PassManager composable, and converting certain class attributes to dynamic properties.

  • Invariant pre- and post-condition checks are now executed within the PassBase call method.
  • PassManager is now a subclass of PassBase with dynamic properties for in_place and destructive attributes.
Comments suppressed due to low confidence (2)

onnxscript/ir/passes/_pass_infra.py:88

  • [nitpick] The precondition error message does not include specific context about which invariant check failed; consider providing additional details to simplify debugging.
raise PreconditionError(

onnxscript/ir/passes/_pass_infra.py:100

  • [nitpick] Include more context in the postcondition error message to better pinpoint the failure, such as details about the pass or model state.
raise PostconditionError(

@justinchuby justinchuby requested a review from Copilot March 22, 2025 15:01
Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR improves the pass infrastructure by running invariant functions at the call method, refactoring PassManager into a composable Pass, and introducing new pass types (Sequential, InPlacePass, OutOfPlacePass, DestructivePass) to clarify pass behavior.

  • In _pass_infra.py, refactors PassBase and adds abstract properties, along with new pass subclasses and additional invariant checks.
  • In init.py, updates the re-exports to include the new pass types.
  • In optimizer modules, updates several passes to inherit from InPlacePass instead of PassBase.

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file
File Description
onnxscript/ir/passes/_pass_infra.py Refactored PassBase, added new pass subclasses with invariant checks, and updated PassManager behavior.
onnxscript/ir/passes/init.py Updated module exports for the new pass classes.
onnxscript/optimizer/_remove_unused_function.py Changed RemoveUnusedFunctionPass to extend InPlacePass.
onnxscript/optimizer/_remove_unused.py Changed RemoveUnusedNodesPass to extend InPlacePass.
onnxscript/optimizer/_constant_folding.py Changed FoldConstantsPass to extend InPlacePass.
Comments suppressed due to low confidence (2)

onnxscript/ir/passes/_pass_infra.py:197

  • The Sequential pass assumes that self.passes is non-empty, which may lead to an IndexError if no passes are provided. Consider adding a check or raising a clear error when self.passes is empty.
return self.passes[0].changes_input or self.passes[0].in_place

onnxscript/ir/passes/_pass_infra.py:202

  • [nitpick] Consider using an explicit identifier (such as pass_.class.name) when logging the pass to improve clarity in debugging output.
logger.debug("Running the %s-th pass '%s'", i, pass_)

return False


class DestructivePass(PassBase):
Copy link
Collaborator

@gramalingam gramalingam Mar 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we really need this? Is there any good example where this is useful? Seems to me it is better to encourage users to use just one of above two ... (which may mean we don't need these two boolean properties) ... on a different (but possibly related) note, I can see the value of "analysis" passes that compute some information but don't modify the model, but not sure whether those are within the scope of "passes" as designed ...

Copy link
Collaborator Author

@justinchuby justinchuby Mar 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#2119 this is an example where the pass is DestructivePass. Maybe possible to implement it as a copy as well? I can remove the class to discourage the use though.

(but the boolean properties are still useful. For example, we can turn a destructive pass to an out of place pass by chaining a "Copy" pass with the destructive pass automatically by looking at these two properties)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#2119 this is an example where the pass is DestructivePass. Maybe possible to implement it as a copy as well? I can remove the class to discourage the use though.

(but the boolean properties are still useful. For example, we can turn a destructive pass to an out of place pass by chaining a "Copy" pass with the destructive pass automatically by looking at these two properties)

I see. I think I would classify #2219 as an example of an in-place pass that returns some extra information (the extracted sub-model) ... and a version that returns only the extracted sub-model without modifying the input-model could be an out-of-place pass (though, in general, we might need the ability to return values other than a model for more complex analysis passes).

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That’s a good point. I will change it to inplace

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Although: the two properties are still useful I suppose: When we compose passes with into a sequence, whether a pass is out of place or destructive requires some careful logic: https://github.com/microsoft/onnxscript/pull/2120/files#diff-70c7e5b3422f4daaf1611d4f76578c96e4c5894cced3d51718efa0290219f7f5R180-R186

@justinchuby justinchuby disabled auto-merge March 24, 2025 21:50
self._in_place = all(pass_.in_place for pass_ in passes)
# The reason changes_inputs is decided by the first pass is that if the first pass is either in-place,
# or if it is not designed to be in-place but somehow changes the input (destructive),
# this pass sequence will change inputs.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I mean would second or other passes that changes inputs after the first pass?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the first pass is functional, the second pass will take the new model that the first pass returns, which means the second pass has no chance to affect the input model.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But if the first pass is side-effect only pass. Shouldn't we check the following passes?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. In that case we just assume it changes the model for now. I think that's ok?

@@ -68,14 +71,72 @@ class PassResult:
class PassBase(abc.ABC):
"""Base class for all passes.

Class attributes:
in_place: Whether the pass modifies the model in place.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to add a paragraph in the documentation of the exporter (torch.onnx.export) to mention the list of passes applied to the fx graph?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure!

@justinchuby
Copy link
Collaborator Author

I will merge since this is the foundation for some other PRs. Please feel free to comment and I will address in a follow up

@justinchuby justinchuby enabled auto-merge (squash) March 26, 2025 14:55
@justinchuby justinchuby merged commit 7d800b6 into main Mar 26, 2025
25 of 29 checks passed
@justinchuby justinchuby deleted the justinchu/pass-infra branch March 26, 2025 14:56
bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025
1. Run invariant functions `requires` and `ensures` by default at Pass
`__call__` to match pytorch's pass behavior. This means the invariants
cannot be too expensive because they are always checked.
2. Make PassManager a `Pass` so that it can be composed.
3. Add `changes_input` attribute to indicate if the input is changed.
Turn two class attributes into properties for them to be dynamic.
Combining the two attributes we can tell if a pass is destructive. For
now the properties are unused but they will become useful when we want
to have a better guard on pass usage etc.
4. Create `Sequential`, `InPlacePass`, `FunctionalPass` to help users
create passes.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025
1. Run invariant functions `requires` and `ensures` by default at Pass
`__call__` to match pytorch's pass behavior. This means the invariants
cannot be too expensive because they are always checked.
2. Make PassManager a `Pass` so that it can be composed.
3. Add `changes_input` attribute to indicate if the input is changed.
Turn two class attributes into properties for them to be dynamic.
Combining the two attributes we can tell if a pass is destructive. For
now the properties are unused but they will become useful when we want
to have a better guard on pass usage etc.
4. Create `Sequential`, `InPlacePass`, `FunctionalPass` to help users
create passes.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
bmehta001 pushed a commit to bmehta001/onnxscript that referenced this pull request Apr 11, 2025
1. Run invariant functions `requires` and `ensures` by default at Pass
`__call__` to match pytorch's pass behavior. This means the invariants
cannot be too expensive because they are always checked.
2. Make PassManager a `Pass` so that it can be composed.
3. Add `changes_input` attribute to indicate if the input is changed.
Turn two class attributes into properties for them to be dynamic.
Combining the two attributes we can tell if a pass is destructive. For
now the properties are unused but they will become useful when we want
to have a better guard on pass usage etc.
4. Create `Sequential`, `InPlacePass`, `FunctionalPass` to help users
create passes.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
module: IR Intermediate representation topic: passes
Projects
Development

Successfully merging this pull request may close these issues.

5 participants