Add fix for writing to closures (pytorch#233)

* Add fix for writing to closures

* run black

* one more time

Co-authored-by: Elias Ellison <[email protected]>

Author: eellison

tests/test_misc.py

@@ -947,6 +947,38 @@ def fn3():
         self.assertEqual(cnts.frame_count, 2)
         self.assertEqual(cnts.op_count, 11)
 
+    def test_write_to_closures_in_inlining(self):
+        out = []
+        for use_dynamo in [False, True]:
+
+            def make_counter():
+                x = torch.randn(10)
+
+                def counter():
+                    nonlocal x
+                    x = x + 1
+                    return x
+
+                return counter
+
+            torch.manual_seed(0)
+            counter = make_counter()
+            if not use_dynamo:
+                out.append(counter() + counter())
+            else:
+                cnts = torchdynamo.testing.CompileCounter()
+
+                @torchdynamo.optimize(cnts, nopython=True)
+                def fn(counter):
+                    return counter() + counter()
+
+                out.append(fn(counter))
+                self.assertEqual(cnts.frame_count, 1)
+                self.assertEqual(cnts.op_count, 3)
+                self.assertFalse(same(counter() + counter(), out[-1]))
+
+        self.assertTrue(same(out[0], out[1]))
+
     def test_top_package_import(self):
         def fn(x):
             import torch.fx

torchdynamo/side_effects.py

@@ -187,6 +187,14 @@ def track_cell_new(
         self.keepalive.append(obj)
         return variable
 
+    def track_cell_existing(self, source: Source, item: Any):
+        variable = variables.NewCellVariable(
+            mutable_local=AttributeMutationExisting(source),
+        )
+        self.id_to_variable[id(item)] = variable
+        self.keepalive.append(item)
+        return variable
+
     def prune_dead_object_new(self, tx):
         live_new_objects = set()
         skip_obj = None
@@ -232,13 +240,14 @@ def codegen(self, cg: PyCodegen):
         ]
 
         for var in modified_vars:
-            if isinstance(var.mutable_local, AttributeMutationNew) and isinstance(
-                var, variables.NewCellVariable
-            ):
+            if isinstance(
+                var.mutable_local, (AttributeMutationExisting, AttributeMutationNew)
+            ) and isinstance(var, variables.NewCellVariable):
                 cg.load_import_from(utils.__name__, "make_cell")
                 cg.extend_output([create_instruction("CALL_FUNCTION", 0)])
                 cg.add_cache(var)
-                var.mutable_local.source = LocalSource(cg.tempvars[var])
+                if isinstance(var.mutable_local, AttributeMutationNew):
+                    var.mutable_local.source = LocalSource(cg.tempvars[var])
             elif isinstance(var.mutable_local, AttributeMutationNew):
                 cg.load_import_from(utils.__name__, "object_new")
                 cg(var.mutable_local.cls_source)

torchdynamo/symbolic_convert.py

@@ -1332,7 +1332,15 @@ def STORE_DEREF(self, inst):
             else:
                 self.output.side_effects.store_cell(cell, val)
         else:
-            unimplemented("write to __closure__ while inlining")
+            if isinstance(
+                self.symbolic_locals.get(inst.argval),
+                torchdynamo.variables.NewCellVariable,
+            ):
+                self.output.side_effects.store_cell(
+                    self.symbolic_locals[inst.argval], self.pop()
+                )
+            else:
+                unimplemented("write to __closure__ while inlining")
 
     def LOAD_DEREF(self, inst):
         if inst.argval in self.closure_cells:
@@ -1342,7 +1350,11 @@ def LOAD_DEREF(self, inst):
             else:
                 self.push(self.output.side_effects.load_cell(cell))
         else:
-            super().LOAD_DEREF(inst)
+            maybe_sym_local = self.symbolic_locals.get(inst.argval, None)
+            if isinstance(maybe_sym_local, torchdynamo.variables.NewCellVariable):
+                self.push(self.output.side_effects.load_cell(maybe_sym_local))
+            else:
+                super().LOAD_DEREF(inst)
 
     def LOAD_CLOSURE(self, inst):
         assert inst.argval in self.cell_and_freevars()

torchdynamo/variables/functions.py

@@ -140,11 +140,29 @@ def bind_args(self, parent, args, kwargs):
                 elif self.source:
                     from .builder import VariableBuilder
 
-                    source = AttrSource(
-                        GetItemSource(AttrSource(self.source, "__closure__"), idx),
-                        "cell_contents",
-                    )
-                    result[name] = VariableBuilder(parent, source)(cell.cell_contents)
+                    side_effects = parent.output.side_effects
+                    if cell in side_effects:
+                        out = side_effects[cell]
+                    else:
+                        closure_cell = GetItemSource(
+                            AttrSource(self.source, "__closure__"), idx
+                        )
+                        closure_cell_contents = AttrSource(
+                            closure_cell, "cell_contents"
+                        )
+
+                        # cells are written to with "cell_contents",
+                        # so the source should just be the closure_cell, not its contents
+                        out = side_effects.track_cell_existing(closure_cell, cell)
+                        side_effects.store_cell(
+                            out,
+                            VariableBuilder(parent, closure_cell_contents)(
+                                cell.cell_contents
+                            ),
+                        )
+
+                    result[name] = out
+
                 else:
                     unimplemented("inline with __closure__")