pytorch · StrongerXi · Feb 11, 2025 · Feb 11, 2025 · Feb 12, 2025 · Feb 13, 2025
diff --git a/torch/_dynamo/output_graph.py b/torch/_dynamo/output_graph.py
@@ -1542,15 +1542,7 @@ def dedup_pass(self):
             return dict()
 
     def install_subgraph(self, name, sub_gm):
-        next_name = None
-        i = 0
-        while not next_name:
-            candidate = f"{name}_{i}"
-            if candidate in self.nn_modules:
-                i += 1
-            else:
-                next_name = candidate
-
+        next_name = get_unique_name_wrt(name, self.nn_modules, requires_suffix=True)
         sub_gm.__name__ = next_name
         sub_gm.torchdynamo_force_dynamic = False
         # This graph module is not present in the user space, so it can't be
@@ -2256,14 +2248,7 @@ def create_graph_input(
                     TracingContext.extract_stack()
                 )
 
-        # unique
-        if name in self.input_name_to_proxy:
-            for i in itertools.count():
-                candidate_name = f"{name}_{i}"
-                if candidate_name not in self.input_name_to_proxy:
-                    name = candidate_name
-                    break
-
+        name = get_unique_name_wrt(name, self.input_name_to_proxy)
         if self.input_name_to_proxy:
             prev_name = next(reversed(self.input_name_to_proxy))
             node = self.input_name_to_proxy[prev_name].node

diff --git a/torch/_dynamo/utils.py b/torch/_dynamo/utils.py
@@ -2574,23 +2574,25 @@ def get_safe_global_name(tx, root, obj):
     return f"{root}_{id(obj)}_c{tx.output.compile_id}"
 
 
-def get_unique_name_wrt(prefix: str, *containers) -> str:
+def is_in(item: Any, *containers) -> bool:
+    for container in containers:
+        if item in container:
+            return True
+    return False
+
+
+def get_unique_name_wrt(prefix: str, *containers, requires_suffix=False) -> str:
     """
     Return a name that starts with `prefix` and is not in any of the
     `containers` (e.g., map, set).
     """
-    name = prefix
+    if not requires_suffix and not is_in(prefix, *containers):
+        return prefix
+
     for i in itertools.count():
-        found = False
-        for container in containers:
-            if name in container:
-                found = True
-                break
-
-        if not found:
-            return name
-        # else update and retry
-        name = f"{prefix}_{i}"
+        candidate = f"{prefix}_{i}"
+        if not is_in(candidate, *containers):
+            return candidate
 
     raise AssertionError("unreachable")