frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-06 00:54:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[dynamo][cpp-guards] Optimize tensor.grad accessor (#123226)

Browse Source 
For LayoutLM model, reduces C++ guard overhead by 1.48x. These are the numbers

![image](https://github.com/pytorch/pytorch/assets/13822661/25cfc35b-b67d-4903-8403-71fa931dacdd)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/123226
Approved by: https://github.com/jansel
This commit is contained in:
Animesh Jain
2024-04-02 16:59:04 -07:00
committed by PyTorch MergeBot
parent 9288b27461
commit d91db70295
4 changed files with 91 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
torch/_dynamo/source.py
View File

@ -170,6 +170,25 @@ class AttrSource(ChainedSource):
return f"{self.base.name()}.{self.member}"
# Represents tensor.grad source. It could be represented by AttrSource as well.
# But, we could access grad field on tensor directly in C++ without going
# through the Python bytecodes. Therefore, we use a separate source for grad
# field.
@dataclasses.dataclass(frozen=True)
class GradSource(ChainedSource):
member: str = "grad"
def reconstruct(self, codegen):
self.base.reconstruct(codegen)
codegen.extend_output(codegen.create_load_attrs(self.member))
def guard_source(self):
return self.base.guard_source()
def name(self):
return f"{self.base.name()}.{self.member}"
@dataclasses.dataclass(frozen=True)
class ParamBufferSource(AttrSource):
def guard_source(self):