CARVIEW |
Navigation Menu
-
Notifications
You must be signed in to change notification settings - Fork 24.7k
sparse_broadcast_to: less memory footprint, fewer kernel launches #142364
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
đź”— Helpful Linksđź§Ş See artifacts and rendered test results at hud.pytorch.org/pr/142364
Note: Links to docs will display an error until the docs builds have been completed. âś… No FailuresAs of commit 1c535d6 with merge base c29b4ed ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
9db7b0a
to
fae893d
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice! This looks good, minor nit in the CHECK message.
aten/src/ATen/native/TensorShape.cpp
Outdated
|
||
for (int64_t i = 0; i < self.dim(); ++i) { | ||
TORCH_CHECK(self_size[i] == 1 || self_size[i] == size[i + new_sparse_dims], | ||
"The input's lenght ", self_size[i], " at dimension ", i, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
"The input's lenght ", self_size[i], " at dimension ", i, | |
"The input's length ", self_size[i], " at dimension ", i, |
aten/src/ATen/native/TensorShape.cpp
Outdated
for (int64_t i = 0; i < self.dim(); ++i) { | ||
TORCH_CHECK(self_size[i] == 1 || self_size[i] == size[i + new_sparse_dims], | ||
"The input's lenght ", self_size[i], " at dimension ", i, | ||
" does not broadcast over the requested shape of lenght ", size[i + new_sparse_dims], |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
" does not broadcast over the requested shape of lenght ", size[i + new_sparse_dims], | |
" does not broadcast over the requested shape of length ", size[i + new_sparse_dims], |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Awesome!
fae893d
to
1c535d6
Compare
Thank you for your reviews, guys! @pytorchbot merge |
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 |
As per title.
The following implementation removes the usage of
repeat_interleave, tile
andfull_coo_indices
and replaces them with broadcasting. That way we reduce memory traffic (and are likely to hit cache a lot) and the total number of launched kernels.cc @alexsamardzic @pearu @cpuhrsch @amjames @bhosmer @jcaip