Trace through pytree API with dynamo. (#108533)

ezyang · pytorchmergebot · commit 518308a740b7 · 2023-09-20T00:04:56.000Z
Fix: #107315 This PR enables dynamo to trace through the `pytree` API by inlining its functions. In order to do so, a few details of `pytree` had to be changed. In summary, this PR: - Introduces `TreeSpecVariable` for representing `TreeSpec` instances - Specializes `<type>.__bases__` call, returning a `TupleVariable` - Enables the call to `id` builtin function for every variable that implements `as_python_constant` method - Specializes `ConstantVariable.call_method` for its (un)flatten functions - Implements `UserDefinedObjectVariable.as_python_constant` - Modifies `pytree` by: - Make `SUPPORTED_NODES` a map of ids (instead of types) to `NodeDef` - Removed `functools.wraps` function, since it can't be inlined Pull Request resolved: #108533 Approved by: https://github.com/ezyang, https://github.com/voznesenskym ghstack dependencies: #109201
diff --git a/test/dynamo/test_export.py b/test/dynamo/test_export.py
@@ -2815,7 +2815,9 @@ def f(x):
             a = A()
             return x.sum() + type(a).func().sum()
 
-        with self.assertRaisesRegex(torch._dynamo.exc.UserError, "Can't call type()"):
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.UserError, r"Can't access members of type\(obj\)"
+        ):
             gm, _ = torch._dynamo.export(f, aten_graph=True)(torch.ones(6, 4))
 
         def f_correct(x):
@@ -3763,7 +3765,6 @@ def forward(self, x):
         self.assertTrue(torch.allclose(gm(inp_test)[1][0], gm2(inp_test)[1][0]))
         self.assertTrue(torch.allclose(gm(inp_test)[1][1], gm2(inp_test)[1][1]))
 
-    @unittest.expectedFailure
     def test_fx_pytree(self):
         def foo(args):
             flat_args, spec = torch.utils._pytree.tree_flatten(args)
diff --git a/test/dynamo/test_higher_order_ops.py b/test/dynamo/test_higher_order_ops.py
@@ -2669,6 +2669,7 @@ def fn(x, y):
         )
         self.assertEqual(actual, expected)
 
+    @unittest.expectedFailure
     def test_vmap_pytree_inputs(self):
         counters.clear()
         x = torch.ones(2, 3)
diff --git a/test/test_pytree.py b/test/test_pytree.py
@@ -265,10 +265,10 @@ def test_treespec_repr_dynamo(self):
         self.assertExpectedInline(
             repr(spec),
             """\
-TreeSpec(TupleVariable, None, [*,
-  TreeSpec(ListVariable, None, [*,
+TreeSpec(tuple, None, [*,
+  TreeSpec(list, None, [*,
     *,
-    TreeSpec(ListVariable, None, [*])])])""",
+    TreeSpec(list, None, [*])])])""",
         )
 
     def test_broadcast_to_and_flatten(self):
diff --git a/torch/_dynamo/allowed_functions.py b/torch/_dynamo/allowed_functions.py
@@ -163,6 +163,7 @@ def _is_allowed_module_prefix(obj):
         disallowed_modules = (
             "torch.optim.",
             "torch.utils._foreach_utils",  # omit the period so we match all the functions in this module
+            "torch.utils._pytree",
             "torch.nn.modules.rnn.",
             "torch._dynamo.",
             "torch._C._dynamo.",
diff --git a/torch/_dynamo/eval_frame.py b/torch/_dynamo/eval_frame.py
@@ -1464,10 +1464,6 @@ def patch():
             # disable future hooking
             opt.step.hooked = True
 
-        torch._dynamo.variables.lists._register_dynamo_list_to_tree_spec()
-        torch._dynamo.variables.lists._register_dynamo_tuple_to_tree_spec()
-        torch._dynamo.variables.dicts._register_dynamo_dict_to_tree_spec()
-
     @staticmethod
     def suppress_torch_distributed_warnings(fn):
         def inner_fn(*args, **kwargs):
diff --git a/torch/_dynamo/guards.py b/torch/_dynamo/guards.py
@@ -71,6 +71,7 @@
 TensorGuards = torch._C._dynamo.guards.TensorGuards
 check_obj_id = torch._C._dynamo.guards.check_obj_id
 check_type_id = torch._C._dynamo.guards.check_type_id
+dict_version = torch._C._dynamo.guards.dict_version
 
 
 # For user stack printing
@@ -105,6 +106,7 @@ def uninteresting_files():
         ("___odict_getitem", collections.OrderedDict.__getitem__),
         ("___dict_param_key_ids", dict_param_key_ids),
         ("___dict_const_keys", dict_const_keys),
+        ("___dict_version", dict_version),
         ("___tuple_iterator_len", tuple_iterator_len),
         ("___tuple_iterator_getitem", tuple_iterator_getitem),
         ("__math_isnan", math.isnan),
@@ -276,6 +278,13 @@ def TYPE_MATCH(self, guard: Guard):
         code = f"___check_type_id({self.arg_ref(guard)}, {obj_id})"
         self._produce_guard_code(guard, [code])
 
+    def DICT_VERSION(self, guard: Guard):
+        # ___check_dict_version is same as `dict_version(x) == y`
+        ref = self.arg_ref(guard)
+        version = dict_version(self.get(guard.name))
+        code = f"___dict_version({ref}) == {version}"
+        self._produce_guard_code(guard, [code])
+
     def BOOL_FALSE(self, guard: Guard):
         # Guard on the runtime value being 'False',
         # can be faster than seemingly equivalent checks like DICT_KEYS for empty dict
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -391,6 +391,19 @@ def _wrap(self, value):
         elif is_namedtuple(value):
             return self.wrap_listlike(value)
 
+        elif value is torch.utils._pytree.SUPPORTED_NODES:
+            result = {
+                k: UserDefinedObjectVariable(
+                    value[k],
+                    source=GetItemSource(self.get_source(), k),
+                    # For SUPPORTED_NODES, we guard on the dictionary version (PEP509)
+                    # under the assumption that the values themselves don't change.
+                    guards=self.make_guards(GuardBuilder.DICT_VERSION),
+                )
+                for k in value.keys()
+            }
+            return ConstDictVariable(result, type(value))
+
         elif istype(
             value, (dict, collections.defaultdict, collections.OrderedDict)
         ) and all(
diff --git a/torch/_dynamo/variables/builtin.py b/torch/_dynamo/variables/builtin.py
@@ -21,7 +21,13 @@
 )
 from ..guards import GuardBuilder
 from ..replay_record import DummyModule
-from ..source import AttrSource, is_constant_source, SuperSource, TypeSource
+from ..source import (
+    AttrSource,
+    GetItemSource,
+    is_constant_source,
+    SuperSource,
+    TypeSource,
+)
 from ..utils import (
     build_checkpoint_variable,
     check_constant_args,
@@ -1060,7 +1066,7 @@ def call_getattr(
             TorchVariable,
             UserFunctionVariable,
         )
-        from .builder import VariableBuilder
+        from .builder import SourcelessBuilder, VariableBuilder
 
         options = VariableTracker.propagate(self, obj, name_var)
         guards = options["guards"]
@@ -1089,6 +1095,27 @@ def call_getattr(
         else:
             source = None
 
+        if name == "__bases__":
+            try:
+                value = obj.as_python_constant()
+                if isinstance(value, type):
+                    bases = value.__bases__
+                    if source is not None:
+                        tuple_args = [
+                            VariableBuilder(tx, GetItemSource(source, i))(b)
+                            for i, b in enumerate(bases)
+                        ]
+                    elif len(bases) == 1 and (
+                        bases[0] is object or bases[0] is torch._C._TensorBase
+                    ):
+                        tuple_args = [SourcelessBuilder()(tx, bases[0])]
+                    else:
+                        unimplemented(f"unexpected sourceless type bases: {bases}")
+
+                    return variables.TupleVariable(tuple_args, **options)
+            except NotImplementedError:
+                pass
+
         if isinstance(obj, variables.NNModuleVariable):
             return obj.var_getattr(tx, name).add_options(options)
         elif isinstance(obj, variables.TensorVariable) and name == "grad":
@@ -1229,6 +1256,9 @@ def call_type(self, tx, obj: VariableTracker):
                 self, obj
             )
 
+        if py_type is not None:
+            return ConstantVariable(py_type)
+
         raise UserError(
             UserErrorType.ANTI_PATTERN,
             f"Can't call type() on generated custom object {obj}. "
diff --git a/torch/_dynamo/variables/constant.py b/torch/_dynamo/variables/constant.py
@@ -4,7 +4,7 @@
 import torch
 
 from .. import variables
-from ..exc import unimplemented
+from ..exc import unimplemented, UserError, UserErrorType
 from ..utils import istype, np
 from .base import typestr, VariableTracker
 
@@ -81,6 +81,12 @@ def unpack_var_sequence(self, tx):
             raise NotImplementedError from e
 
     def const_getattr(self, tx, name):
+        if isinstance(self.value, type):
+            raise UserError(
+                UserErrorType.ANTI_PATTERN,
+                "Can't access members of type(obj) for a generated custom object. "
+                "Please use __class__ instead",
+            )
         member = getattr(self.value, name)
         if callable(member):
             raise NotImplementedError()
diff --git a/torch/_dynamo/variables/dicts.py b/torch/_dynamo/variables/dicts.py
@@ -2,13 +2,10 @@
 import dataclasses
 import functools
 import inspect
-from typing import Any, Dict, List, Tuple
+from typing import Dict, List
 
 import torch
 import torch.fx
-from torch.fx import _pytree as fx_pytree
-
-from torch.utils import _pytree as pytree
 
 from .. import variables
 from ..bytecode_transformation import create_call_function, create_instruction
@@ -586,40 +583,3 @@ def var_getattr(self, tx, name: str) -> "VariableTracker":
 
     def call_hasattr(self, tx, name: str) -> "VariableTracker":
         return variables.ConstantVariable(hasattr(self.obj, name)).add_options(self)
-
-
-def _dictvariable_flatten(d: ConstDictVariable) -> Tuple[List[Any], pytree.Context]:
-    if d.python_type() is not dict:
-        # Note - ConstDictVariable can contain different kinds of dicts.
-        # However, flattening for those must differ and so cannot share the same registration as even if we
-        # consult the underlying python_type() to guide our flattening, that data will need to be propagated
-        # to unflatten. We do not have a good mechanism of doing this today, so we find it easier to treat this
-        # as unimplemented for now.
-
-        # TODO - Add support for flattening a ConstDictVariable with any underlying user_cls
-        unimplemented(f"Unsupported flattening of {d.python_type()}")
-    return list(d.items.values()), list(d.items.keys())
-
-
-def _dictvariable_unflatten(
-    values: List[Any], context: pytree.Context
-) -> ConstDictVariable:
-    assert all(isinstance(x, VariableTracker) for x in values)
-
-    # Guard propagation happens in the ConstDictVariable constructor
-    return ConstDictVariable(
-        dict(zip(context, values)), user_cls=dict, mutable_local=MutableLocal()
-    )
-
-
-def _register_dynamo_dict_to_tree_spec():
-    pytree._register_pytree_node(
-        ConstDictVariable,
-        _dictvariable_flatten,
-        _dictvariable_unflatten,
-    )
-
-    fx_pytree.register_pytree_flatten_spec(
-        ConstDictVariable,
-        _dictvariable_flatten,
-    )
diff --git a/torch/_dynamo/variables/lists.py b/torch/_dynamo/variables/lists.py
@@ -3,12 +3,10 @@
 import functools
 import inspect
 import operator
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import torch
 import torch.fx
-from torch.fx import _pytree as fx_pytree
-from torch.utils import _pytree as pytree
 
 from .. import variables
 from ..bytecode_transformation import create_call_function, create_instruction
@@ -182,25 +180,6 @@ def list_eq(tx, left, right):
             comps,
         ).add_options(options)
 
-    # List-like implementations for pytree
-    def __len__(self):
-        return len(self.items)
-
-    def __getitem__(self, index):
-        if isinstance(index, slice):
-            index = SliceVariable(
-                [
-                    ConstantVariable(index.start),
-                    ConstantVariable(index.stop),
-                    ConstantVariable(index.step),
-                ]
-            )
-        elif isinstance(index, int):
-            index = ConstantVariable(index)
-        else:
-            raise TypeError("Invalid index type. Must be int or slice.")
-        return self.getitem_const(index)
-
 
 class RangeVariable(BaseListVariable):
     def __init__(self, items, **kwargs):
@@ -732,56 +711,6 @@ class TupleIteratorVariable(ListIteratorVariable):
     pass
 
 
-def _listvariable_flatten(d: ListVariable) -> Tuple[List[Any], pytree.Context]:
-    return d.items, None
-
-
-def _listvariable_unflatten(values: List[Any], context: pytree.Context) -> ListVariable:
-    assert all(isinstance(x, VariableTracker) for x in values)
-
-    # Guard propagation happens in the BaseListVariable constructor
-    return ListVariable(values, mutable_local=MutableLocal())
-
-
-def _register_dynamo_list_to_tree_spec():
-    pytree._register_pytree_node(
-        ListVariable,
-        _listvariable_flatten,
-        _listvariable_unflatten,
-    )
-
-    fx_pytree.register_pytree_flatten_spec(
-        ListVariable,
-        _listvariable_flatten,
-    )
-
-
-def _tuplevariable_flatten(d: TupleVariable) -> Tuple[List[Any], pytree.Context]:
-    return d.items, None
-
-
-def _tuplevariable_unflatten(
-    values: List[Any], context: pytree.Context
-) -> TupleVariable:
-    assert all(isinstance(x, VariableTracker) for x in values)
-
-    # Guard propagation happens in the BaseListVariable constructor
-    return TupleVariable(values)
-
-
-def _register_dynamo_tuple_to_tree_spec():
-    pytree._register_pytree_node(
-        TupleVariable,
-        _tuplevariable_flatten,
-        _tuplevariable_unflatten,
-    )
-
-    fx_pytree.register_pytree_flatten_spec(
-        TupleVariable,
-        _tuplevariable_flatten,
-    )
-
-
 class SetVariable(VariableTracker):
     @dataclasses.dataclass
     class SetElement:
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
diff --git a/torch/_export/wrappers.py b/torch/_export/wrappers.py
diff --git a/torch/csrc/dynamo/guards.cpp b/torch/csrc/dynamo/guards.cpp
diff --git a/torch/utils/_pytree.py b/torch/utils/_pytree.py
