frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Stop accessing func._schema in _python_dispatch.correct_storage_aliasing (#161292)

Browse Source 
func._schema is a pybind, accessing the arguments/returns is expensive, we have no reason to do it anyway, and even though #161301 makes accessing the arguments/returns less expensive, this still seems to improve performance.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/161292
Approved by: https://github.com/wconstab, https://github.com/malfet, https://github.com/bdhirsh
ghstack dependencies: #161301
This commit is contained in:
Scott Wolchok
2025-08-29 15:54:11 -07:00
committed by PyTorch MergeBot
parent 5d35b49ba7
commit 1a64bf2636
1 changed files with 3 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
torch/utils/_python_dispatch.py
View File

@ -561,13 +561,10 @@ and output of type {type(ret)}. But expected types to match."""
assert isinstance(ret, torch.Tensor), f"type: {type(ret)}"
torch._functionalize_unsafe_set(ret, arg)
num_args = len(func._schema.arguments)
num_returns = len(func._schema.returns)
for arg_idx in range(num_args):
for return_idx in range(num_returns):
schema_arg = schema_info.args[arg_idx]
for arg_idx, schema_arg in enumerate(schema_info.args):
for return_idx, schema_out in enumerate(schema_info.outs):
is_read_only_alias_match = (
schema_arg.alias_set & schema_info.outs[return_idx].alias_set
schema_arg.alias_set & schema_out.alias_set
) and not schema_arg.is_write
if is_read_only_alias_match:
alias_non_inplace_storage(args[arg_idx], outs[return_idx])