8000 cache dict and string rep for better perf (#158372) · pytorch/pytorch@0cb36e2 · GitHub
[go: up one dir, main page]

Skip to content

Commit 0cb36e2

Browse files
Gasoonjiapytorchmergebot
authored andcommitted
cache dict and string rep for better perf (#158372)
Summary: NodeSouce should not be updated after created, so that it would be better if we cache its dict and string representation for better perf. Test Plan: ci Rollback Plan: Reviewed By: yushangdi Differential Revision: D78298501 Pull Request resolved: #158372 Approved by: https://github.com/yushangdi
1 parent 584a051 commit 0cb36e2

File tree

1 file changed

+20
-11
lines changed

1 file changed

+20
-11
lines changed

torch/fx/traceback.py

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,10 @@ def __init__(
8080
self.node_info = None
8181
self.from_node = []
8282

83+
# cache the action string and dict representation for performance.
84+
self._action_string = None
85+
self._dict = None
86+
8387
@property
8488
def name(self) -> str:
8589
return self.node_info.name if self.node_info else ""
@@ -96,7 +100,9 @@ def __repr__(self):
96100
return self.print_readable()
97101

98102
def _get_action_string(self):
99-
return "+".join([a.name.lower() for a in self.action])
103+
if self._action_string is None:
104+
self._action_string = "+".join([a.name.lower() for a in self.action])
105+
return self._action_string
100106

101107
def print_readable(self, indent=0):
102108
if indent > 9:
@@ -112,16 +118,19 @@ def print_readable(self, indent=0):
112118
return result
113119

114120
def to_dict(self) -> dict:
115-
# Convert the object to a dictionary
116-
action_string = self._get_action_string()
117-
return {
118-
"name": self.name,
119-
"target": self.target,
120-
"graph_id": self.graph_id,
121-
"pass_name": self.pass_name,
122-
"action": action_string,
123-
"from_node": [node.to_dict() for node in self.from_node],
124-
}
121+
if self._dict is None:
122+
# Convert the object to a dictionary
123+
action_string = self._get_action_string()
124+
self._dict = {
125+
"name": self.name,
126+
"target": self.target,
127+
"graph_id": self.graph_id,
128+
"pass_name": self.pass_name,
129+
"action": action_string,
130+
"from_node": [node.to_dict() for node in self.from_node],
131+
}
132+
133+
return self._dict
125134

126135
def __eq__(self, other: object):
127136
if not isinstance(other, NodeSource):

0 commit comments

Comments
 (0)
0