diff --git a/torch/autograd/profiler.py b/torch/autograd/profiler.py index fa43af2701171..e757714636b8e 100644 --- a/torch/autograd/profiler.py +++ b/torch/autograd/profiler.py @@ -1072,14 +1072,20 @@ class EnforceUnique: def __init__(self): self.seen = set() + self._cache = {} def see(self, *key): r""" Observe a key and raise an error if it is seen multiple times. """ + if key in self._cache: + return self._cache[key] + if key in self.seen: raise RuntimeError("duplicate key: " + str(key)) self.seen.add(key) + self._cache[key] = True + return True def parse_nvprof_trace(path): @@ -1193,7 +1199,8 @@ class KinetoStepTracker: """ _current_step = 0 - _step_dict: dict[str, int] = defaultdict(int) + _step_dict: dict[str, int] = {} + _cache: dict[str, int] = {} @classmethod def init_step_count(cls, requester: str): @@ -1207,7 +1214,8 @@ def erase_step_count(cls, requester: str) -> bool: r""" Remove a given requester. """ - return cls._step_dict.pop(requester, None) is not None + result = cls._step_dict.pop(requester, None) is not None + return result @classmethod def increment_step(cls, requester: str) -> int: @@ -1216,11 +1224,14 @@ def increment_step(cls, requester: str) -> int: Additionally if the max over all step counts has incremented then trigger the _kineto_step() returns global step count """ + if requester in cls._cache: + return cls._cache[requester] + if requester not in cls._step_dict: cls.init_step_count(requester) cls._step_dict[requester] += 1 - new_step = max(cls._step_dict.values()) + new_step = max(cls._step_dict.values()) if cls._step_dict else 0 if new_step > cls._current_step: delta = new_step - cls._current_step if delta > 1: @@ -1232,6 +1243,8 @@ def increment_step(cls, requester: str) -> int: for _ in range(delta): _kineto_step() cls._current_step = new_step + + cls._cache[requester] = cls._current_step return cls._current_step @classmethod