-
Notifications
You must be signed in to change notification settings - Fork 64
[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
[IR] Improve pass infra #2120
Conversation
There was a problem hiding this 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.
❌ 4 Tests Failed:
View the top 3 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
There was a problem hiding this 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.
There was a problem hiding this 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()."
There was a problem hiding this 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(
There was a problem hiding this 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_)
onnxscript/ir/passes/_pass_infra.py
Outdated
return False | ||
|
||
|
||
class DestructivePass(PassBase): |
There was a problem hiding this comment.
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 ...
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done
There was a problem hiding this comment.
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).
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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
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. |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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. | |||
|
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure!
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 |
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>
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>
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>
requires
andensures
by default at Pass__call__
to match pytorch's pass behavior. This means the invariants cannot be too expensive because they are always checked.Pass
so that it can be composed.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.Sequential
,InPlacePass
,FunctionalPass
to help users create passes.