Cache optimization

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