Introduce repeat and RepeatOp #3687
Merged
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
Almost same semantics PyTorch repeat. Previously only partially introduced as a translation from a repeat pattern using concat. which had a bug when repeating broadcast IDs. This PR fixes the issue by handling broadcast separately using a new IR node, RepeatOp, which represents repetition of broadcast IDs.
|
!test |
wujingyue
reviewed
Jan 9, 2025
| } | ||
|
|
||
| std::string RepeatOp::toInlineString(int indent_size) const { | ||
| NVF_CHECK(false, "Tensor op can not be printed inline"); |
There was a problem hiding this comment.
Question for my understanding: is it correct that only IterDomain ops can be printed inline according to our convention?
There was a problem hiding this comment.
Looks like that's the case, but I just copied this from ExpandOp::toInlineString.
| sizes.reserve(out()->getLogicalDomain().size()); | ||
| const auto c2p = | ||
| PairwiseLogicalDomainMap(in(), out()).mapConsumerToProducer(); | ||
| for (const auto i : c10::irange(out()->getLogicalDomain().size())) { |
There was a problem hiding this comment.
Since this is the output, it should not have a reduction ID. I added an assertion in the constructor.
| @@ -504,6 +504,10 @@ class ComputeAtLogicalDomainMapBuilder : private BackwardVisitor { | |||
| mapPointwiseLikeOp(op); | |||
| } | |||
|
|
|||
| void handle(RepeatOp* op) override { | |||
| mapPointwiseLikeOp(op); | |||
There was a problem hiding this comment.
Is it supposed to be mapped as pointwise? The input and the output don't even have the same extent.
wujingyue
approved these changes
Jan 9, 2025
|
!build |
|
!build |
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.
Adds
repeatas an alias op as well as theRepeatOpIR node. Therepeatop has almost the same semantics as the PyTorch repeat.The main motivation is to fix #3682, which is due to #3645, which introduced a preseg pass that detects and translates a repeat pattern to broadcast, expand and reshape. The issue of #3682 is because that the translation-based method does not work when a broadcast ID is repeated. I originally just used
TensorDomain::flatten(https://github.com/NVIDIA/Fuser/blob/main/csrc/ir/nodes.cpp#L3674-L3740), which just merges broadcast IDs. However, for reshape, it should not merge but squeeze them. Merging broadcast IDs triggered an assertion of the transpose scheduler as seen in #3682.TensorDomain::flattenneeds to be fixed (#3691), but that's a separate issue. For fixing #3682, since repeating broadcast IDs cannot be translated to the broadcast-expand-reshape pattern anyway, I added the newRepeatOpnode. I initially thought it could be justLoadStoreOpbut decided to have a different IR node since, unlike usual LoadStore case, some of the broadcast IDs of a producer becomes concrete IDs in the corresponding consumer logical domain. I did actually try usingLoadStoreOpbut some of the preseg passes complained the mismatched broadcast pattern.Repeating non-broadcast IDs is still done by the broadcast-expand-reshape patten. Only for repeating broadcast IDs gets represented using the
RepeatOpnode.Fixes #3682