[dynamo] Refactor OrderedDict to dict (pytorch#113234)

In Python3 all dicts are ordered.

Pull Request resolved: pytorch#113234
Approved by: https://github.com/oulgen, https://github.com/lezcano

Author: jansel

Diff for: torch/_dynamo/codegen.py

@@ -3,7 +3,7 @@
 import re
 import sys
 import types
-from typing import Counter, List, Optional, OrderedDict
+from typing import Counter, Dict, List, Optional
 
 import torch.nn
 from . import utils
@@ -51,9 +51,7 @@ def __init__(
         self.root = root
         self.top_of_stack: Optional[VariableTracker] = None
         self.uses: Counter[VariableTracker] = collections.Counter()
-        self.graph_outputs: OrderedDict[
-            int, GraphOutputEntry
-        ] = collections.OrderedDict()
+        self.graph_outputs: Dict[int, GraphOutputEntry] = {}
         self._output: List[Instruction] = []
         self.tempvars = tempvars or {}
         self.tx = tx

Diff for: torch/_dynamo/convert_frame.py

@@ -405,11 +405,9 @@ def maybe_cprofile(func):
     return func
 
 
-from collections import OrderedDict
-
 from torch.utils.hooks import RemovableHandle
 
-_bytecode_hooks: Dict[int, BytecodeHook] = OrderedDict()
+_bytecode_hooks: Dict[int, BytecodeHook] = {}
 
 
 def register_bytecode_hook(hook: BytecodeHook) -> RemovableHandle:

Diff for: torch/_dynamo/guards.py

@@ -91,46 +91,37 @@ def uninteresting_files():
     return {inspect.getfile(m) for m in mods}
 
 
-CLOSURE_VARS = collections.OrderedDict(
-    [
-        ("___check_type_id", check_type_id),
-        ("___check_obj_id", check_obj_id),
-        (
-            "___current_backend",
-            lambda: torch._dynamo.eval_frame.guarded_backend_cache.current_backend,
-        ),
-        (
-            "___lookup_backend",
-            lambda backend_obj_id: torch._dynamo.eval_frame.guarded_backend_cache.cached_backends[
-                backend_obj_id
-            ],
-        ),
-        (
-            "___skip_backend_check",
-            lambda: torch._dynamo.eval_frame.guarded_backend_cache.skip_backend_check_for_run_only_mode,
-        ),
-        ("___odict_getitem", collections.OrderedDict.__getitem__),
-        ("___dict_param_key_ids", dict_param_key_ids),
-        ("___dict_const_keys", dict_const_keys),
-        ("___dict_version", dict_version),
-        ("___dict_contains", lambda a, b: a in b),
-        ("___tuple_iterator_len", tuple_iterator_len),
-        ("___tuple_iterator_getitem", tuple_iterator_getitem),
-        ("__math_isnan", math.isnan),
-        ("inf", float("inf")),
-        ("__load_module", lambda name: importlib.import_module(name)),
-        ("utils_device", torch.utils._device),
-        ("device", torch.device),
-        (
-            "___from_numpy",
-            # If not numpy array, piggy back on e.g. tensor guards to check type
-            lambda a: torch.as_tensor(a)
-            if isinstance(a, (np.generic, np.ndarray))
-            else a,
-        ),
-        ("torch", torch),
-    ]
-)
+CLOSURE_VARS = {
+    "___check_type_id": check_type_id,
+    "___check_obj_id": check_obj_id,
+    "___current_backend": (
+        lambda: torch._dynamo.eval_frame.guarded_backend_cache.current_backend
+    ),
+    "___lookup_backend": (
+        lambda backend_obj_id: torch._dynamo.eval_frame.guarded_backend_cache.cached_backends[
+            backend_obj_id
+        ]
+    ),
+    "___skip_backend_check": (
+        lambda: torch._dynamo.eval_frame.guarded_backend_cache.skip_backend_check_for_run_only_mode
+    ),
+    "___odict_getitem": collections.OrderedDict.__getitem__,
+    "___dict_param_key_ids": dict_param_key_ids,
+    "___dict_const_keys": dict_const_keys,
+    "___dict_version": dict_version,
+    "___dict_contains": lambda a, b: a in b,
+    "___tuple_iterator_len": tuple_iterator_len,
+    "___tuple_iterator_getitem": tuple_iterator_getitem,
+    "__math_isnan": math.isnan,
+    "inf": float("inf"),
+    "__load_module": lambda name: importlib.import_module(name),
+    "utils_device": torch.utils._device,
+    "device": torch.device,
+    "___from_numpy":
+    # If not numpy array, piggy back on e.g. tensor guards to check type
+    (lambda a: torch.as_tensor(a) if isinstance(a, (np.generic, np.ndarray)) else a),
+    "torch": torch,
+}
 
 if sys.version_info[:2] <= (3, 8):
     # [Note: Python Version <= 3.8]
@@ -1138,17 +1129,15 @@ def convert(size_or_stride):
         if global_state is None:
             # we should only hit this case in NopTests()
             global_state = convert_frame.GlobalStateGuard()
-        closure_vars = collections.OrderedDict(
-            [
-                ("___guarded_code", self),
-                ("___check_tensors", check_tensors_fn),
-                ("___check_tensors_verbose", check_tensors_verbose_fn),
-                ("___check_global_state", global_state.check),
-                ("tensor_check_names", tensor_check_names),
-            ]
-            + list(SYMPY_INTERP.items())
-        )
-        closure_vars.update(CLOSURE_VARS)
+        closure_vars = {
+            "___guarded_code": self,
+            "___check_tensors": check_tensors_fn,
+            "___check_tensors_verbose": check_tensors_verbose_fn,
+            "___check_global_state": global_state.check,
+            "tensor_check_names": tensor_check_names,
+            **SYMPY_INTERP,
+            **CLOSURE_VARS,
+        }
 
         unique_code_parts = list(unique(code_parts))
         make_guard_fn_args = ", ".join(closure_vars.keys())

Diff for: torch/_dynamo/output_graph.py

@@ -10,18 +10,7 @@
 import traceback
 import weakref
 from dataclasses import dataclass
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    NamedTuple,
-    Optional,
-    OrderedDict,
-    Set,
-    Tuple,
-    Union,
-)
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Set, Tuple, Union
 
 import sympy
 
@@ -827,9 +816,7 @@ def append_prefix_insts():
         root = FakeRootModule(self.nn_modules)
         # Add all the local vars to the "stack" so restore at the end
         restore_vars = []
-        val_to_names: OrderedDict[
-            VariableTracker, List[str]
-        ] = collections.OrderedDict()
+        val_to_names: Dict[VariableTracker, List[str]] = {}
         if stack_values:
             val_to_names[stack_values[-1]] = list()
         # NB: Typically (i.e., for graph compile from RETURN_VALUE),
@@ -1307,7 +1294,7 @@ def __init__(
         # Map from graph input name to its placeholder proxy object, where the
         # map's keys give all current placeholder node names and can be used to
         # create unique node names
-        self.input_name_to_proxy: OrderedDict[str, fx.Proxy] = collections.OrderedDict()
+        self.input_name_to_proxy: Dict[str, fx.Proxy] = {}
         # Node => computed real value (see utils.get_real_value)
         self.real_value_cache: Dict[fx.Node, torch.Tensor] = {}
 
@@ -1324,9 +1311,8 @@ def __init__(
         # - If we are tracing a HigherOrderOperator's body_fn, then we
         # need to keep track of what free variables were lifted so we can
         # rewrite the HigherOrderOperator call using the traced body_fn.
-        # This is a OrderedDict so that we can
-        # maintain the order of args for the HigherOrderOperator call.
-        self.lifted_freevars = collections.OrderedDict()
+        # Dicts maintain the order of args for the HigherOrderOperator call.
+        self.lifted_freevars = {}
         self.prev_inst = None
 
         self._cur_code = None

Diff for: torch/_dynamo/side_effects.py

@@ -1,4 +1,3 @@
-import collections
 import inspect
 from typing import Any, Dict, List, Optional
 
@@ -76,8 +75,8 @@ def __init__(
         tensor_hooks=None,
     ):
         super().__init__()
-        self.id_to_variable = id_to_variable or collections.OrderedDict()
-        self.store_attr_mutations = store_attr_mutations or collections.OrderedDict()
+        self.id_to_variable = id_to_variable or {}
+        self.store_attr_mutations = store_attr_mutations or {}
         self.keepalive = keepalive or []
         self.save_for_backward = save_for_backward or []
         self.tensor_hooks = tensor_hooks or {}
@@ -115,11 +114,10 @@ def diff(self, other: "SideEffects") -> Optional[str]:
     def clone(self):
         """Create a shallow copy"""
         return self.__class__(
-            id_to_variable=collections.OrderedDict(self.id_to_variable),
-            store_attr_mutations=collections.OrderedDict(
-                (k, collections.OrderedDict(v))
-                for k, v in self.store_attr_mutations.items()
-            ),
+            id_to_variable=dict(self.id_to_variable),
+            store_attr_mutations={
+                k: dict(v) for k, v in self.store_attr_mutations.items()
+            },
             keepalive=list(self.keepalive),
             save_for_backward=self.save_for_backward,
             tensor_hooks=self.tensor_hooks,
@@ -129,14 +127,14 @@ def apply(self, fn, cache=None, skip_fn=lambda _: False):
         if cache is None:
             cache = dict()
 
-        self.id_to_variable = collections.OrderedDict(
-            (k, VariableTracker.apply(fn, v, cache, skip_fn))
+        self.id_to_variable = {
+            k: VariableTracker.apply(fn, v, cache, skip_fn)
             for k, v in self.id_to_variable.items()
-        )
-        self.store_attr_mutations = collections.OrderedDict(
-            (k, VariableTracker.apply(fn, v, cache, skip_fn))
+        }
+        self.store_attr_mutations = {
+            k: VariableTracker.apply(fn, v, cache, skip_fn)
             for k, v in self.store_attr_mutations.items()
-        )
+        }
         self.save_for_backward = VariableTracker.apply(
             fn, self.save_for_backward, cache, skip_fn
         )
@@ -164,7 +162,7 @@ def store_attr(self, item: VariableTracker, name: str, value: VariableTracker):
         assert self.is_attribute_mutation(item)
         self.check_allowed_side_effect(item)
         if item.mutable_local not in self.store_attr_mutations:
-            self.store_attr_mutations[item.mutable_local] = collections.OrderedDict()
+            self.store_attr_mutations[item.mutable_local] = {}
         self.store_attr_mutations[item.mutable_local][name] = value
 
     def load_attr(self, item, name, deleted_ok=False):
@@ -320,12 +318,12 @@ def is_live(var: VariableTracker):
         for skip_obj, setattrs in self.store_attr_mutations.items():
             VariableTracker.apply(visit, setattrs)
 
-        self.id_to_variable = collections.OrderedDict(
-            (k, v) for k, v in self.id_to_variable.items() if is_live(v)
-        )
-        self.store_attr_mutations = collections.OrderedDict(
-            (k, v) for k, v in self.store_attr_mutations.items() if is_live(k)
-        )
+        self.id_to_variable = {
+            k: v for k, v in self.id_to_variable.items() if is_live(v)
+        }
+        self.store_attr_mutations = {
+            k: v for k, v in self.store_attr_mutations.items() if is_live(k)
+        }
 
     def mutation(self, oldvar, newvar):
         self.check_allowed_side_effect(oldvar)

Diff for: torch/_dynamo/symbolic_convert.py

@@ -622,9 +622,9 @@ def prune_dead_locals(self):
         # reads = reads | {"__class__"}
         # output variables?
         reads = reads | set(self.cell_and_freevars())
-        self.symbolic_locals = collections.OrderedDict(
-            [(k, v) for k, v in self.symbolic_locals.items() if k in reads]
-        )
+        self.symbolic_locals = {
+            k: v for k, v in self.symbolic_locals.items() if k in reads
+        }
         self.output.side_effects.prune_dead_object_new(self)
 
     def call_function(
@@ -1863,7 +1863,7 @@ def copy_graphstate(self) -> InstructionTranslatorGraphState:
         """Create a checkpoint of the current state by copying everything"""
         return InstructionTranslatorGraphState(
             self.output.copy_graphstate(),
-            collections.OrderedDict(self.symbolic_locals),
+            dict(self.symbolic_locals),
             list(self.stack),
             list(self.block_stack),
             self.instruction_pointer,
@@ -2071,9 +2071,9 @@ def __init__(
             f_globals=f_globals,
             f_builtins=f_builtins,
             code_options=code_options,
-            symbolic_locals=collections.OrderedDict(),  # set below
+            symbolic_locals={},  # set below
             # A global var is inserted only after a STORE_GLOBAL happens to it
-            symbolic_globals=collections.OrderedDict(),
+            symbolic_globals={},
             f_code=f_code,
             export=export,
             inline_depth=0,

Diff for: torch/_dynamo/types.py

@@ -1,17 +1,7 @@
 import dataclasses
 import sys
 import types
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    NamedTuple,
-    Optional,
-    OrderedDict,
-    Protocol,
-    Union,
-)
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Protocol, Union
 
 from typing_extensions import TypeAlias
 
@@ -41,7 +31,7 @@ class GuardFail(NamedTuple):
 
 
 class GuardFn(Protocol):
-    closure_vars: OrderedDict[str, object]
+    closure_vars: Dict[str, object]
     args: List[str]
     code_parts: List[str]
     verbose_code_parts: List[str]

Diff for: torch/_dynamo/utils.py

@@ -80,10 +80,10 @@
 log = logging.getLogger(__name__)
 
 # profiling compilation time by function
-compilation_time_metrics = collections.OrderedDict()
+compilation_time_metrics = {}
 
 # profiling compilation time by frame phase
-frame_phase_timing = collections.OrderedDict()
+frame_phase_timing = {}
 
 timer_counter = itertools.count()
 

Diff for: torch/_dynamo/variables/base.py

@@ -212,11 +212,7 @@ def update_object_dict(v):
             result = [cls.apply(fn, v, cache, skip_fn) for v in value]
         elif istype(value, tuple):
             result = tuple(cls.apply(fn, v, cache, skip_fn) for v in value)
-        elif istype(value, collections.OrderedDict):
-            result = collections.OrderedDict(
-                cls.apply(fn, v, cache, skip_fn) for v in value.items()
-            )
-        elif istype(value, dict):
+        elif istype(value, (dict, collections.OrderedDict)):
             assert "__name__" not in value, "_nonvar_fields should have excluded this"
             result = {
                 k: cls.apply(fn, v, cache, skip_fn) for k, v in list(value.items())