Prevent cudaStreamSync when indexing GPU tensors with boolean CPU mask #156384
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/156384
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit c172e4e with merge base 6b7767f ( UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "release notes: cuda" |
mikaylagawarecki
added
the
triaged
|
@pytorchmergebot merge |
|
Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge. Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt. If you think this is a mistake, please contact PyTorch Dev Infra. |
|
dfdf9bd should fix the build failure on CI |
|
@ngimel any chance you could trigger CI to verify that this PR is good to go? |
|
@pytorchmergebot merge -r |
|
@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here |
|
Successfully rebased |
pytorchmergebot
force-pushed
the
index-put-no-sync
branch
from
dfdf9bd to
c172e4e
Compare
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 |
6 participants
index_putwith a boolean mask (target[mask] = src) causes acudaStreamSynchronize. When bothmaskandtargettensors are on GPU this is expected.However, the sync can be prevented if the
maskis a CPU tensor.Internally a new index tensor is created with
mask.nonzero()so we can use a non-blocking copy to transfer it to the GPU since it cannot be accidentally mutated by the user between its creation and the device copy. @ngimel Let me know if I'm missing something.I think this is useful since users can't prevent a sync simply by making sure all tensors are on the same device as with other ops. Instead one would need to do something like this which is much less readable
Fixes #12461