Remove arith dependency from DynamicIotaOpToBroadcast#2918
Open
brucechanglongxu wants to merge 1 commit intoopenxla:mainfrom
Open
Remove arith dependency from DynamicIotaOpToBroadcast#2918brucechanglongxu wants to merge 1 commit intoopenxla:mainfrom
brucechanglongxu wants to merge 1 commit intoopenxla:mainfrom
Conversation
The pattern was using arith.index_cast to convert index-typed shapes to i64, pulling in an unnecessary arith dialect dependency. Now it bails out on index-typed shapes (users should run shape-legalize-to-stablehlo first) and uses stablehlo.convert for non-i64 integer types. Also skips the convert entirely when the shape is already i64. This removes the arith::ArithDialect dependency from the StablehloAggressiveSimplification pass. Fixes openxla#2628
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
1 participant
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.
The DynamicIotaOpToBroadcast pattern was pulling in an arith dialect
dependency to handle index-typed shapes via arith.index_cast. This is
unnecessary — by the time aggressive simplification runs, shapes should
already be legalized to integer types via shape-legalize-to-stablehlo.
Now the pattern bails out on index-typed shapes instead of introducing
arith ops, and the ArithDialect dependency is dropped from the pass
entirely. For non-i64 integer shapes (e.g. i32), it still uses
stablehlo.convert; for i64 it skips the conversion altogether.
Updated the test to verify the index case is left untouched.
Fixes #2628