Introduce torch.sym_add, variadic add #138660
Closed
Conversation
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
Partially addresses #128150 When you have big sums of values, we end up computing long chains of binary addition in our FX graph representation. Not only is this ugly, it also is quadratic, as the sympy.Add constructor is O(N) in number of arguments. Instead, ensure that we maintain the summation as a single FX node so we can do the entire addition all in one go. Signed-off-by: Edward Z. Yang <[email protected]> [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138660
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ You can merge normally! (1 Unrelated Failure)As of commit c84b2af with merge base 72dde6e ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/inductor
module: cpu
ezyang
approved these changes
Oct 23, 2024
Contributor
Author
|
@pytorchbot merge -i |
Collaborator
Collaborator
Tested internally here: https://www.internalfb.com/diff/D64057744 This is a reland after previous internal failures. main change is ``` if min is None and max is None: torch._check_is_size(size) return ``` Partially addresses #128150 When you have big sums of values, we end up computing long chains of binary addition in our FX graph representation. Not only is this ugly, it also is quadratic, as the sympy.Add constructor is O(N) in number of arguments. Instead, ensure that we maintain the summation as a single FX node so we can do the entire addition all in one go. Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc jgong5 mingfeima XiaobingSuper sanchitintel ashokei jingxu10 ezyang SherlockNoMad EikanWang wenzhe-nrv voznesenskym penguinwu Guobing-Chen zhuhaozhe blzheng jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov rec [ghstack-poisoned]
laithsakka
added a commit
that referenced
this pull request
Oct 23, 2024
Partially addresses #128150 When you have big sums of values, we end up computing long chains of binary addition in our FX graph representation. Not only is this ugly, it also is quadratic, as the sympy.Add constructor is O(N) in number of arguments. Instead, ensure that we maintain the summation as a single FX node so we can do the entire addition all in one go. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: a170651 Pull Request resolved: #138660
Contributor
Author
|
rebase |
XuehaiPan
reviewed
Oct 23, 2024
Comment on lines
+1016
to
+1046
| # Special case for sum on tuple/list of ints | ||
| if ( | ||
| self.fn is builtins.sum | ||
| and len(args) == 1 | ||
| and not kwargs | ||
| and isinstance(args[0], (variables.ListVariable, variables.TupleVariable)) | ||
| and all( | ||
| (isinstance(x, variables.ConstantVariable) and isinstance(x.value, int)) | ||
| or (isinstance(x, variables.SymNodeVariable) and x.python_type() is int) | ||
| for x in args[0].items | ||
| ) | ||
| ): | ||
| return variables.SymNodeVariable.create( | ||
| tx, | ||
| tx.output.create_proxy( | ||
| "call_function", | ||
| torch.sym_sum, | ||
| (tuple(a.as_proxy() for a in args[0].items),), | ||
| {}, | ||
| ), | ||
| sym_num=torch.sym_sum( | ||
| [ | ||
| ( | ||
| x.value | ||
| if isinstance(x, variables.ConstantVariable) | ||
| else x.sym_num | ||
| ) | ||
| for x in args[0].items | ||
| ] | ||
| ), | ||
| ) |
Collaborator
There was a problem hiding this comment.
I'd like to create a dispatch registry to make the code easier to maintain.
check_fn, dispatch_fn = self.shortcuts.get(self.fn, (None, None))
if check_fn is not None and check_fn(args, kwargs):
return dispatch_fn(self, tx, args, kwargs)
bobrenjc93
approved these changes
Oct 23, 2024
Contributor
Author
|
@pytorchbot merge |
Collaborator
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
This was referenced Oct 24, 2024
SamGinzburg
pushed a commit
that referenced
this pull request
Oct 28, 2024
Tested internally here: https://www.internalfb.com/diff/D64057744 This is a reland after previous internal failures. main change is ``` if min is None and max is None: torch._check_is_size(size) return ``` Partially addresses #128150 When you have big sums of values, we end up computing long chains of binary addition in our FX graph representation. Not only is this ugly, it also is quadratic, as the sympy.Add constructor is O(N) in number of arguments. Instead, ensure that we maintain the summation as a single FX node so we can do the entire addition all in one go. Signed-off-by: Edward Z. Yang <[email protected]> Pull Request resolved: #138660 Approved by: https://github.com/ezyang, https://github.com/bobrenjc93
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.
Stack from ghstack (oldest at bottom):
Tested internally here: https://www.internalfb.com/diff/D64057744
This is a reland after previous internal failures.
main change is
Partially addresses #128150
When you have big sums of values, we end up computing long chains of
binary addition in our FX graph representation. Not only is this ugly,
it also is quadratic, as the sympy.Add constructor is O(N) in number
of arguments. Instead, ensure that we maintain the summation as a
single FX node so we can do the entire addition all in one go.
Signed-off-by: Edward Z. Yang [email protected]
cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @ezyang @SherlockNoMad @EikanWang @wenzhe-nrv @voznesenskym @penguinwu @Guobing-Chen @zhuhaozhe @blzheng @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @rec