[dynamo] Validate mutation_type and source in VariableTracker.__init__ (pytorch#141717)

As title, this also uncovered a few invalid use cases; the cases that
cause error are fixed in separate patches prior to this patch, and the
rest are fixed in this patch.

This patch also moves a few `.source` mutation to variable construction,
to increase the coverage of the validation.

Fixes pytorch#133027.

Pull Request resolved: pytorch#141717
Approved by: https://github.com/jansel
ghstack dependencies: pytorch#141713, pytorch#141714, pytorch#141715, pytorch#141902, pytorch#141716

Author: StrongerXi

torch/_dynamo/variables/base.py

@@ -479,6 +479,23 @@ def __init__(
         self.source = source
         self.mutation_type = mutation_type
 
+        # NOTE sometimes mutation_type is set afterwards for implementation
+        # convenience, we don't validate those cases at the moment.
+        if mutation_type is not None:
+            if isinstance(mutation_type, (ValueMutationNew, AttributeMutationNew)):
+                # If this fails, it's either
+                # 1. one mistakenly passed in a source
+                # 2. `mutation_type` is incorrect
+                assert source is None
+            else:
+                assert isinstance(
+                    mutation_type, (ValueMutationExisting, AttributeMutationExisting)
+                )
+                # If this fails, it's either
+                # 1. one forgot to pass in a source
+                # 2. `mutation_type` is incorrect
+                assert source is not None
+
 
 def typestr(*objs):
     if len(objs) == 1:

torch/_dynamo/variables/builder.py

@@ -431,11 +431,6 @@ def install_guards(self, *guards):
         install_guard(*[source.make_guard(guard) for guard in guards], skip=1)
         return {}
 
-    def set_source_and_track_mutable(self, value, var):
-        assert isinstance(var, VariableTracker)
-        var.source = self.source
-        return self.tx.output.side_effects.track_mutable(value, var)
-
     @classmethod
     def _type_dispatch(cls):
         return cls._type_dispatch_impl(config.trace_numpy)
@@ -607,7 +602,6 @@ def create_2d_tma_descriptor():
         elif CustomizedDictVariable.is_matching_cls_hf(type(value)):
             self.install_guards(GuardBuilder.TYPE_MATCH)
             result = CustomizedDictVariable.wrap(self, value)
-            result.source = self.source
             return self.tx.output.side_effects.track_object_existing(value, result)
         elif istype(value, (dict, collections.defaultdict, collections.OrderedDict)):
             self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
@@ -671,7 +665,7 @@ def build_key_value(i, k, v):
                     result, user_cls=type(value), source=self.source
                 )
 
-            return self.set_source_and_track_mutable(value, result)
+            return self.tx.output.side_effects.track_mutable(value, result)
         elif isinstance(value, torch.nn.Module):
             return self.wrap_module(value)
         elif ConstantVariable.is_literal(value):  # non-atomic literals
@@ -1137,7 +1131,7 @@ def build_key_value(i, k, v):
             )
         elif RestrictedListSubclassVariable.is_matching_cls(type(value)):
             self.install_guards(GuardBuilder.SEQUENCE_LENGTH)
-            return self.set_source_and_track_mutable(
+            return self.tx.output.side_effects.track_mutable(
                 value,
                 RestrictedListSubclassVariable(
                     [
@@ -1148,6 +1142,7 @@ def build_key_value(i, k, v):
                     ],
                     user_cls=type(value),
                     user_cls_source=AttrSource(self.source, "__class__"),
+                    source=self.source,
                 ),
             )
         elif TorchScriptObjectVariable.is_matching_cls(type(value)):
@@ -1326,9 +1321,9 @@ def wrap_listlike(self, value: Union[tuple, list, odict_values, NamedTuple]):
             )
             tensor_list_proxy.node.meta["grapharg"] = grapharg
 
-        result = BaseListVariable.cls_for_instance(value)(output)
+        result = BaseListVariable.cls_for_instance(value)(output, source=self.source)
         if istype(value, (list, collections.deque)):
-            return self.set_source_and_track_mutable(value, result)
+            return self.tx.output.side_effects.track_mutable(value, result)
         return result
 
     def wrap_tuple_iterator(self, value: tuple_iterator):
@@ -1339,11 +1334,8 @@ def wrap_tuple_iterator(self, value: tuple_iterator):
             )
             for i in range(tuple_iterator_len(value))
         ]
-        result = TupleIteratorVariable(
-            output, mutation_type=ValueMutationNew(), source=self.source
-        )
-
-        return self.set_source_and_track_mutable(value, result)
+        result = TupleIteratorVariable(output, source=self.source)
+        return self.tx.output.side_effects.track_mutable(value, result)
 
     def wrap_range_iterator(self, value: range_iterator):
         self.install_guards(GuardBuilder.RANGE_ITERATOR_MATCH)
@@ -1512,7 +1504,7 @@ def wrap_literal(self, value):
             self.install_guards(GuardBuilder.CONSTANT_MATCH)
             result = ConstantVariable.create(value=value, source=self.source)
             if isinstance(value, (list, set)):
-                return self.set_source_and_track_mutable(value, result)
+                return self.tx.output.side_effects.track_mutable(value, result)
             return result
 
     def assert_not_wrapped_by_this_graph(self, value: torch.Tensor):
@@ -2403,7 +2395,7 @@ def handle_traced_output(example_value, tx, proxy, options, subclass_type, targe
         elif istype(example_value, tuple):
             return TupleVariable(unpacked, **options)
         elif istype(example_value, (list, immutable_list)):
-            return ListVariable(unpacked, mutation_type=ValueMutationNew(), **options)
+            return ListVariable(unpacked, **options)
         else:
             assert example_value.__class__.__module__ == "torch.return_types" or hasattr(
                 example_value, "_fields"

torch/_dynamo/variables/builtin.py

@@ -1375,7 +1375,9 @@ def call_custom_dict(tx: "InstructionTranslator", user_cls, *args, **kwargs):
                     arg, user_cls, mutation_type=ValueMutationNew()
                 )
             elif isinstance(arg, variables.ConstDictVariable):
-                return arg.clone(user_cls=user_cls, mutation_type=ValueMutationNew())
+                return arg.clone(
+                    user_cls=user_cls, source=None, mutation_type=ValueMutationNew()
+                )
             elif isinstance(
                 arg,
                 (

torch/_dynamo/variables/dicts.py

@@ -845,7 +845,7 @@ def wrap(cls, builder, obj):
                 if val is not None:
                     key = ConstantVariable.create(key)
                     items[key] = var
-        return cls(items, user_cls)
+        return cls(items, user_cls, source=builder.source)
 
     def __init__(self, items, user_cls, **options) -> None:
         super().__init__(items, user_cls, **options)

torch/_dynamo/variables/misc.py

@@ -721,11 +721,11 @@ def visit(node):
 
     def call_backward(self, tx: "InstructionTranslator", args, kwargs):
         fn = self.fn_cls.backward
-        self.source = AttrSource(self.source, "backward")
         assert type(args[0].value) is torch._dynamo.external_utils.FakeBackwardCFunction
         assert isinstance(fn, types.FunctionType)
 
-        return variables.UserFunctionVariable(fn, source=self.source).call_function(
+        fn_source = AttrSource(self.source, "backward")
+        return variables.UserFunctionVariable(fn, source=fn_source).call_function(
             tx, args, kwargs
         )
 

torch/_dynamo/variables/torch.py

@@ -1117,6 +1117,9 @@ def call_nn_parameter(cls, tx, data=None, requires_grad=True):
                 (data.as_proxy(), placeholder.as_proxy()),
                 {},
             ),
+            # In reconstruct() we should use the original parameter. The one
+            # returned by the graph will be an alias.
+            source=placeholder.source,
         )
         assert isinstance(result, variables.TensorVariable)
         result.class_type = torch.nn.Parameter
@@ -1127,9 +1130,6 @@ def call_nn_parameter(cls, tx, data=None, requires_grad=True):
         # has_grad_fn field to False to workaround the issue.
         result.has_grad_fn = False
 
-        # In reconstruct() should use the original parameter.  The one returned by the graph will be an alias.
-        result.source = placeholder.source
-
         # TODO(jansel): if the new param falls out of scope, currently it won't get freed until
         # the end of the graph.  We should fix this.
         return result