simplify_reshapes: skip find_reshape_dot when it would change element…#4994
Draft
ycastill2-amd wants to merge 2 commits into
Draft
simplify_reshapes: skip find_reshape_dot when it would change element…#4994ycastill2-amd wants to merge 2 commits into
ycastill2-amd wants to merge 2 commits into
Conversation
… count find_reshape_dot moves a reshape across a dot and then reshapes the new dot's output back to the original dot's shape. When the rewrite changes the total element count (e.g. the second input's reshape preserves its contraction axis but changes N), that trailing reshape is invalid and compilation aborts with "Reshape: Wrong number of elements" (observed on an RT-DETR detection model on gfx1150). Skip the rewrite, and drop the speculatively-created dot, when the new dot's element count differs from the original dot's. Signed-off-by: Yviel Castillejos <ycastill@amd.com>
Contributor
|
Thank you for your contribution! Since this is an external pull request, a maintainer must review PR and add the "ok-to-test" label if it is approved for testing. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #4994 +/- ##
===========================================
- Coverage 92.73% 92.71% -0.03%
===========================================
Files 594 596 +2
Lines 31340 31494 +154
===========================================
+ Hits 29063 29197 +134
- Misses 2277 2297 +20
🚀 New features to boost your workflow:
|
Collaborator
|
Can you add a unit test for this? |
Add reshape_dot_changed_element_count, which builds a reshape -> dot -> reshape pattern where the second input's reshape preserves its contraction axis but changes N. Moving the reshapes across the dot would change the total element count and make the trailing reshape invalid, so find_reshape_dot must skip the rewrite and leave the module unchanged. Signed-off-by: Yviel Castillejos <ycastill@amd.com>
| // Moving the reshape across the dot is only valid when the element count is | ||
| // preserved; otherwise the trailing reshape back to the original dot shape is | ||
| // invalid (observed creating an N->M element reshape that aborts compilation). | ||
| if(new_dot->get_shape().elements() != dot->get_shape().elements()) |
Collaborator
There was a problem hiding this comment.
I think we should do this before we add the new instruction to avoid doing a remove_instruction below. Just have an early return. The number of elements can be inferred from the inp and new_other for the check.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens the simplify_reshapes pass by preventing find_reshape_dot from performing a reshape-across-dot rewrite when it would change the dot output’s total element count, which would make the trailing reshape invalid and trigger a reshape element-count error during compilation.
Changes:
- Add an element-count guard in
find_reshape_dotand bail out when the rewrite would change the dot output size. - Add a regression test ensuring the pass leaves the module unchanged for a reshape-dot case where the rewrite would change element count.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/simplify_reshapes.cpp |
Adds element-count validation to skip the reshape-dot rewrite when it would make the trailing reshape invalid. |
test/simplify_reshapes_test.cpp |
Adds a regression test covering the element-count-changing reshape-dot scenario. |
Comment on lines
+1774
to
+1781
| // Moving the reshape across the dot is only valid when the element count is | ||
| // preserved; otherwise the trailing reshape back to the original dot shape is | ||
| // invalid (observed creating an N->M element reshape that aborts compilation). | ||
| if(new_dot->get_shape().elements() != dot->get_shape().elements()) | ||
| { | ||
| m.remove_instruction(new_dot); | ||
| return; | ||
| } |
Comment on lines
+4554
to
+4559
| // find_reshape_dot would move the reshapes across the dot and then reshape the new | ||
| // dot's output back to the original dot's shape. Here the second input's reshape | ||
| // preserves its contraction axis (second-to-last dim) but changes the free (N) dim, | ||
| // so the rewrite changes the dot's total element count and that trailing reshape is | ||
| // invalid (the "Reshape: Wrong number of elements" abort observed on an RT-DETR | ||
| // detection model on gfx1150). The pass must leave the module unchanged. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
… count
find_reshape_dot moves a reshape across a dot and then reshapes the new dot's output back to the original dot's shape. When the rewrite changes the total element count (e.g. the second input's reshape preserves its contraction axis but changes N), that trailing reshape is invalid and compilation aborts with "Reshape: Wrong number of elements" (observed on an RT-DETR detection model on gfx1150).
Skip the rewrite, and drop the speculatively-created dot, when the new dot's element count differs from the original dot's.
Motivation
Technical Details
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot Applicable