[rewriter] Transpose initializer rule #2255
Conversation
❌ 6 Tests Failed:
View the top 3 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
Pull Request Overview
This PR adds a new rewrite rule to collapse Transpose nodes into initializers for constant folding, which removes the need to modify the optimizer’s size limit when folding constants.
- Introduces the TransposeInitializer rule to handle initializer folding.
- Implements the rewriting and checking logic to ensure the initializer meets usage requirements.
- Registers the new rule in the rewriter’s initializer.
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| onnxscript/rewriter/transpose_initializer.py | Adds a new rewrite rule for folding Transpose nodes into initializers. |
| onnxscript/rewriter/init.py | Registers the new transpose_initializer rule in the rule set. |
Comments suppressed due to low confidence (1)
onnxscript/rewriter/transpose_initializer.py:28
- [nitpick] Consider renaming 'original_transpose' to 'transpose_node' to more clearly indicate that it represents the Transpose node consumer of the initializer.
original_transpose = initializer.consumers()[0]
|
I think this is less efficient than an IR pass. Because it will match all transpose nodes, but we only need to look at users of the initializers. Thoughts? @titaiwangms @shubhambhokare1 @gramalingam |
|
But this is op-specific. I don't think it fits in the common passes. So I prefer that it is a rewrite rule. |
Co-authored-by: Shubham Bhokare <32080845+shubhambhokare1@users.noreply.github.com>
With the check inside the rule, I assume it should not be too bad? |
|
In the original issue, there are also transpose -> transpose and transpose -> MatMul/Gemm pattern matching. Should we have those in this rule as well? If we are going to go through Transpose nods anyway (I forget if we can check whether multiple rules are applicable within the same traverse.). |
|
I think we should have different rules since they will match different nodes |
|
Could we add onnxscript/onnxscript/rewriter/generic_pattern_test.py Lines 530 to 573 in 33f31ca as part of this PR to upstream all the fusions done in https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/fusion_constant_fold.py |
|
@shubhambhokare1 I see it's already in , was it not applied? |
Ah I see, however the one in https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/fusion_constant_fold.py seems to be doing, Whereas the current rewrite rule only replace two transposes with an equivalent single transpose |
Sounds good. I will update |
|
@shubhambhokare1 does this one or this one do it? |
justinchuby
changed the title
| if new_value.const_value is None: | ||
| new_value.const_value = old_value.const_value | ||
| if new_value.name is None: | ||
| new_value.name = old_value.name |
There was a problem hiding this comment.
Maybe create an option? Retain new value shapes?
| if perm_attr is not None: | ||
| perm = perm_attr.as_ints() | ||
| else: | ||
| perm = None |
There was a problem hiding this comment.
Can we eliminate that case in def check?
There was a problem hiding this comment.
When perm is None the transpose can still be evaluated. So I don't think it can be eliminated in check()?
|
Replaced by #2355 |
Create a rule to absorb a Transpose node into the initializer (fold constant). This rule is in place to handle initializers of any size. So users do not need to change the size limit of the optimizer when constant folding.
TODO
Fix #2158