[Codegen] Resolve constantOp with multiple layouts users. #11
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
…out. The motivation behind this PR was to solve the issue where despite us emitting to_layout op at certain locations, we ended up generating redundant to_simd earlier in the graph. This results us to not have control on where the layout conversion actually happens. One example of this case is when we emit a to_layout conversion for FP8 Attention. We emit a to_layout right before the 2nd contract and right after the truncf to FP8. Since we do enforcement(backward propagation) first and then propagate (forward propagation), we ended up generating two layout conversions, one the original one we emitted, the second one determined by layout analysis and placed in the middle of a elemwise op that is part of the softmax. This forces us to do shuffles/layout resolution in FP32 which hurts the performance. Additionally, this is also not where we(compiler writers) intended the layout conversion to happen. To solve this issue, we re-ordered the layout analysis to do forward propagation first and then enforcement(backward propagation). This gives layout set by forward propagation precedence, which makes the to_layout emitting more intuitive and behave more as expected. Additionally, we made modification to getAgreedLayout to handle elemwise propagation where some operands are not expected to have layout such as arith.select who's condition operand probably don't have layouts. Signed-off-by: Stanley Winata <[email protected]>
Signed-off-by: Stanley Winata <[email protected]>
Signed-off-by: Stanley Winata <[email protected]>
Main motivation is to handle distribution of constantOp who has users with different layouts. Original use case is to ensure we can distribute attention when the tile size for M,K1,N is the same. Which means the init of 1st contract, and IV's init uses the same constantOp. Since constantOp can only hold a single layout, but multiple to_layout ops with different layouts, for each user, there will be non resolved to_layout op(s). only one of the to_layout op can be resolved properly, the rest would be a "non trivial" resolution since layouts are different. To solve this issue, we introduce a mechanism that detect these cases and make a copy of the arith.constant that get used by other users, when we are trying to resolve for the current constantOp. Signed-off-by: Stanley Winata <[email protected]>
raikonenfnu
force-pushed
the
fixConstantOpLayout2
branch
from
e321e68 to
361448d
Compare
raikonenfnu
force-pushed
the
anchorToLayoutOp
branch
from
ef2233f to
26c8a44
Compare
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.
Main motivation is to handle distribution of constantOp who has users
with different layouts.
Original use case is to ensure we can distribute attention when the tile
size for M,K1,N is the same. Which means the init of 1st contract, and
IV's init uses the same constantOp.
Since constantOp can only hold a single layout, but multiple to_layout ops
with different layouts, for each user, there will be non resolved to_layout
op(s). only one of the to_layout op can be resolved properly, the rest would
be a "non trivial" resolution since layouts are different.
To solve this issue, we introduce a mechanism that detect these cases
and make a copy of the arith.constant that get used by other users, when
we are trying to resolve for the current constantOp.