Loop Counter Increment in Collective Pipeliner #19363
Open
+247
−1
Conversation
This file contains 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
…n of the collective pipeliner.
|
Thanks. I guess this PR fixes the CP bug. But there is another bug that causes corrupted numerics when the stack layer is unrolled. @hawkinsp |
|
@yliu120 what is the other bug? Is it related to this PR? |
Unrelated. This PR fixed the issue in the CP pass and please go ahead! |
| [](const HloInstruction* instruction) { | ||
| return HloPredicateIsOp<HloOpcode::kWhile>(instruction); | ||
| }); | ||
| EXPECT_EQ(while_count, 1); |
There was a problem hiding this comment.
Can you rewrite this part to filecheck? All you want to check is the increment so it should not be long, I think.
There was a problem hiding this comment.
I removed some of the verifications duplicated in other tests but didn't switch to FileCheck since e.g. TransformIncrementIndexByOneBackwards doesn't use it either. Let me know if you want to use it anyway.
There was a problem hiding this comment.
Without so many assertions I think in practice it's now less mandated :)
|
@philipphack Any chance you can address Greg's comments soon? I'd like to get this PR into the next release of JAX (tomorrow, hopefully). |
…n of the collective pipeliner.
|
I added a set of collective pipeliner tests to collective_ops_e2e_test.cc. |
golechwierowicz
approved these changes
Nov 20, 2024
copybara-service bot
pushed a commit
that referenced
this pull request
Nov 20, 2024
Imported from GitHub PR #19363 Sets the loop iteration counter increment in the backward transformation of the collective pipeliner pass to account for cases with non-zero initial value of the loop iteration counter. See #16953 and #18568. Copybara import of the project: -- 06137aa by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. -- 6da45bc by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. Merging this change closes #19363 FUTURE_COPYBARA_INTEGRATE_REVIEW=#19363 from philipphack:u_pipeliner_increment_xla 6da45bc PiperOrigin-RevId: 698294264
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Nov 20, 2024
Imported from GitHub PR openxla/xla#19363 Sets the loop iteration counter increment in the backward transformation of the collective pipeliner pass to account for cases with non-zero initial value of the loop iteration counter. See #16953 and #18568. Copybara import of the project: -- 06137aa0618d372e2d4badbf16920bead9922cfb by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. -- 6da45bcb26643d8994bf608f05230fa748286b02 by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. Merging this change closes #19363 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#19363 from philipphack:u_pipeliner_increment_xla 6da45bcb26643d8994bf608f05230fa748286b02 PiperOrigin-RevId: 698294264
copybara-service bot
pushed a commit
that referenced
this pull request
Nov 20, 2024
Imported from GitHub PR #19363 Sets the loop iteration counter increment in the backward transformation of the collective pipeliner pass to account for cases with non-zero initial value of the loop iteration counter. See #16953 and #18568. Copybara import of the project: -- 06137aa by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. -- 6da45bc by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. Merging this change closes #19363 FUTURE_COPYBARA_INTEGRATE_REVIEW=#19363 from philipphack:u_pipeliner_increment_xla 6da45bc PiperOrigin-RevId: 698294264
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Nov 20, 2024
Imported from GitHub PR openxla/xla#19363 Sets the loop iteration counter increment in the backward transformation of the collective pipeliner pass to account for cases with non-zero initial value of the loop iteration counter. See #16953 and #18568. Copybara import of the project: -- 06137aa0618d372e2d4badbf16920bead9922cfb by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. -- 6da45bcb26643d8994bf608f05230fa748286b02 by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. Merging this change closes #19363 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#19363 from philipphack:u_pipeliner_increment_xla 6da45bcb26643d8994bf608f05230fa748286b02 PiperOrigin-RevId: 698294264
copybara-service bot
pushed a commit
that referenced
this pull request
Nov 20, 2024
Imported from GitHub PR #19363 Sets the loop iteration counter increment in the backward transformation of the collective pipeliner pass to account for cases with non-zero initial value of the loop iteration counter. See #16953 and #18568. Copybara import of the project: -- 06137aa by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. -- 6da45bc by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. Merging this change closes #19363 FUTURE_COPYBARA_INTEGRATE_REVIEW=#19363 from philipphack:u_pipeliner_increment_xla 6da45bc PiperOrigin-RevId: 698294264
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Nov 20, 2024
Imported from GitHub PR openxla/xla#19363 Sets the loop iteration counter increment in the backward transformation of the collective pipeliner pass to account for cases with non-zero initial value of the loop iteration counter. See #16953 and #18568. Copybara import of the project: -- 06137aa0618d372e2d4badbf16920bead9922cfb by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. -- 6da45bcb26643d8994bf608f05230fa748286b02 by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. Merging this change closes #19363 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#19363 from philipphack:u_pipeliner_increment_xla 6da45bcb26643d8994bf608f05230fa748286b02 PiperOrigin-RevId: 698294264
copybara-service bot
pushed a commit
that referenced
this pull request
Nov 20, 2024
Imported from GitHub PR #19363 Sets the loop iteration counter increment in the backward transformation of the collective pipeliner pass to account for cases with non-zero initial value of the loop iteration counter. See #16953 and #18568. Copybara import of the project: -- 06137aa by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. -- 6da45bc by Philipp Hack <[email protected]>: Modifies the loop counter increment set in the backward transformation of the collective pipeliner. Merging this change closes #19363 FUTURE_COPYBARA_INTEGRATE_REVIEW=#19363 from philipphack:u_pipeliner_increment_xla 6da45bc PiperOrigin-RevId: 698294264
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.
Sets the loop iteration counter increment in the backward transformation of the collective pipeliner pass to account for cases with non-zero initial value of the loop iteration counter. See #16953 and #18568.