Optimize mutable torch.library.custom_op overhead #139513
Conversation
We don't need to do a loop over all the args, kwargs in the AdInplaceOrView key; we just need to bump the version on the args, kwargs that are mutable. On the benchmark mentioned in #139494 this made the time go from ``` mutate2 = 61.72943878173828 no_mutate2 = 36.89440155029297 mutate = 236.3092498779297 no_mutate = 59.31964874267578 ``` to ``` mutate2 = 47.976478576660156 no_mutate2 = 38.37468719482422 mutate = 71.21315002441406 no_mutate = 59.7432975769043 ``` Test Plan: - existing tests [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/139513
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 56c2229 with merge base 5e4c8b6 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
We don't need to do a loop over all the args, kwargs in the AdInplaceOrView key; we just need to bump the version on the args, kwargs that are mutable. On the benchmark mentioned in #139494 this made the time go from ``` mutate2 = 61.72943878173828 no_mutate2 = 36.89440155029297 mutate = 236.3092498779297 no_mutate = 59.31964874267578 ``` to ``` mutate2 = 47.976478576660156 no_mutate2 = 38.37468719482422 mutate = 71.21315002441406 no_mutate = 59.7432975769043 ``` Test Plan: - existing tests ghstack-source-id: 46c3beb Pull Request resolved: #139513
| for idx in mutated_idxs: | ||
| increment_version(args[idx]) | ||
| for key in mutated_keys: | ||
| increment_version(kwargs[key]) |
There was a problem hiding this comment.
if we're worried about python overhead here, increment_version() should (as of recently) support Iterable[Tensor] as an argument: https://github.com/pytorch/pytorch/blob/main/torch/autograd/graph.py#L226
There was a problem hiding this comment.
nice, thanks for pointing that out
There was a problem hiding this comment.
It's not clear to me if building an iterator to pass to increment_version is less expensive than calling increment_version in a loop, but I'll try it out
Makes sense, let me file another issue |
zou3519
added
ciflow/trunk
|
@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 |
We don't need to do a loop over all the args, kwargs in the AdInplaceOrView key; we just need to bump the version on the args, kwargs that are mutable. On the benchmark mentioned in pytorch#139494 this made the time go from ``` mutate2 = 61.72943878173828 no_mutate2 = 36.89440155029297 mutate = 236.3092498779297 no_mutate = 59.31964874267578 ``` to ``` mutate2 = 47.976478576660156 no_mutate2 = 38.37468719482422 mutate = 71.21315002441406 no_mutate = 59.7432975769043 ``` Test Plan: - existing tests Pull Request resolved: pytorch#139513 Approved by: https://github.com/bdhirsh ghstack dependencies: pytorch#139509
3 participants
Stack from ghstack (oldest at bottom):
We don't need to do a loop over all the args, kwargs in the
AdInplaceOrView key; we just need to bump the version on the args,
kwargs that are mutable.
On the benchmark mentioned in
#139494
this made the time go from
to
Test Plan: