Repository: pytorch/torchdynamo

Issue:
Side effect to OrderedDict doesn't allow resumption
```
from typing import List
import torch
import torchdynamo
import collections

torchdynamo.config.debug = True

def toy_example(a, b, d):
    x = a + b
    d[3] = 1
    return x * b

d = collections.OrderedDict()
with torchdynamo.optimize("eager"):
    for i in range(4):
        toy_example(torch.randn(10), torch.randn(10), d)
print(d)
```

prints

```
MODIFIED BYTECODE
  8           0 LOAD_GLOBAL              0 (__compiled_fn_0)
              2 LOAD_FAST                0 (a)
              4 LOAD_FAST                1 (b)
              6 CALL_FUNCTION            2
              8 UNPACK_SEQUENCE          1
             10 STORE_FAST               3 (x)

 10          12 LOAD_CONST               1 (1)
             14 LOAD_FAST                2 (d)
             16 LOAD_CONST               2 (3)
             18 STORE_SUBSCR

 11          20 LOAD_FAST                3 (x)
             22 LOAD_FAST                1 (b)
             24 BINARY_MULTIPLY
             26 RETURN_VALUE
```

I expected there to be generated continuation function so torchdynamo could capture x * b.


Discussion:
Should be somewhat easy to add support for this as [regular dict mutation](https://github.com/facebookresearch/torchdynamo/blob/main/torchdynamo/side_effects.py#L273) is already supported (without a graph break) and [all dicts are ordered in Python 3.7+](https://medium.com/junior-dev/python-dictionaries-are-ordered-now-but-how-and-why-5d5a40ee327f).  So just a matter of making OrderedDict go down the same codepath as dict.
There is also a more general question I have, which is that many opcodes don't have `break_graph_if_unsupported`. In principle, wouldn't it be sound to break graph for all unsupported opcodes, assuming that everything on the stack can be restored?

(That is to say, the resolution you suggest above would let us capture an entire graph. But I'm also interested in why there isn't a continuation, if this were some user op, for example.)
You could get there to be a continuation by:
1) Adding a [break_graph_if_unsupported](https://github.com/facebookresearch/torchdynamo/blob/f2efb835a5092e817b5ab12b8a62ae8fd3fc7e90/torchdynamo/symbolic_convert.py#L135) decorator to [STORE_SUBSCR](https://github.com/facebookresearch/torchdynamo/blob/f2efb835a5092e817b5ab12b8a62ae8fd3fc7e90/torchdynamo/symbolic_convert.py#L638)
2) Changing the 1 to a 0 on [this line](https://github.com/facebookresearch/torchdynamo/blob/f2efb835a5092e817b5ab12b8a62ae8fd3fc7e90/torchdynamo/symbolic_convert.py#L149).  `STORE_SUBSCR` pushes 0 items unto the stack, while currently `break_graph_if_unsupported` assumes 1.  We should move this to an arg.
3) Comment out [this line](https://github.com/facebookresearch/torchdynamo/blob/f2efb835a5092e817b5ab12b8a62ae8fd3fc7e90/torchdynamo/symbolic_convert.py#L151).  Same issue as above, `break_graph_if_unsupported` assumes 1 thing is pushed, which needs to get passed to the continuation function.

In principle adding `break_graph_if_unsupported` to way more places would be sound provided you got the stack_effect right.  It already checks that everything on the stack can be restored.

I'll add that one challenge (that is only partially solved right now) of `resume_at_X` continuations is exception handlers.  This gets mapped to the concept of a "block" in python bytecode.  I think this is solvable, but would require more work.

This example doesn't use try/except though -- so not an issue here.

Resolving patch:
diff --git a/torchdynamo/symbolic_convert.py b/torchdynamo/symbolic_convert.py
--- a/torchdynamo/symbolic_convert.py
+++ b/torchdynamo/symbolic_convert.py
@@ -132,28 +132,31 @@ def inner(self: "InstructionTranslatorBase", inst: Instruction):
     return inner
 
 
-def break_graph_if_unsupported(inner_fn):
-    @functools.wraps(inner_fn)
-    def wrapper(self: "InstructionTranslatorBase", inst: Instruction):
-        state = self.copy_graphstate()
-        try:
-            return inner_fn(self, inst)
-        except Unsupported as exc:
-            if not self.should_compile_partial_graph():
-                raise
-            exc.remove_from_stats()
-            exc.add_to_stats("graph_break")
-        self.restore_graphstate(state)
-        self.output.compile_subgraph(self)
-        # note, assuming inst pushes 1
-        self.popn(1 - dis.stack_effect(inst.opcode, inst.arg))
-        self.output.add_output_instructions([inst])
-        self.push(UnknownVariable())
-        self.output.add_output_instructions(
-            self.create_call_resume_at(self.next_instruction)
-        )
+def break_graph_if_unsupported(*, push):
+    def decorator(inner_fn):
+        @functools.wraps(inner_fn)
+        def wrapper(self: "InstructionTranslatorBase", inst: Instruction):
+            state = self.copy_graphstate()
+            try:
+                return inner_fn(self, inst)
+            except Unsupported as exc:
+                if not self.should_compile_partial_graph():
+                    raise
+                exc.remove_from_stats()
+                exc.add_to_stats("graph_break")
+            self.restore_graphstate(state)
+            self.output.compile_subgraph(self)
+            self.popn(push - dis.stack_effect(inst.opcode, inst.arg))
+            self.output.add_output_instructions([inst])
+            for _ in range(push):
+                self.push(UnknownVariable())
+            self.output.add_output_instructions(
+                self.create_call_resume_at(self.next_instruction)
+            )
+
+        return wrapper
 
-    return wrapper
+    return decorator
 
 
 class InstructionTranslatorBase(object):
@@ -536,13 +539,13 @@ def COMPARE_OP(self, inst):
     def GET_ITER(self, inst):
         self.call_function(BuiltinVariable(iter), [self.pop()], {})
 
-    @break_graph_if_unsupported
+    @break_graph_if_unsupported(push=1)
     def CALL_FUNCTION(self, inst):
         args = self.popn(inst.argval)
         fn = self.pop()
         self.call_function(fn, args, {})
 
-    @break_graph_if_unsupported
+    @break_graph_if_unsupported(push=1)
     def CALL_FUNCTION_EX(self, inst):
         if inst.argval == 0:
             kwargsvars = ConstDictVariable({})
@@ -579,7 +582,7 @@ def CALL_FUNCTION_EX(self, inst):
 
         self.call_function(fn, argsvars.items, kwargsvars.items)
 
-    @break_graph_if_unsupported
+    @break_graph_if_unsupported(push=1)
     def CALL_FUNCTION_KW(self, inst):
         argnames = self.pop()
         args = self.popn(inst.argval)
@@ -635,6 +638,7 @@ def STORE_ATTR(self, inst):
             self.create_call_resume_at(self.next_instruction)
         )
 
+    @break_graph_if_unsupported(push=0)
     def STORE_SUBSCR(self, inst):
         val, obj, key = self.popn(3)
         result = obj.call_method(self, "__setitem__", [key, val], {})
@@ -898,7 +902,7 @@ def BUILD_STRING(self, inst):
     BINARY_MODULO = stack_op(operator.mod)
     BINARY_ADD = stack_op(operator.add)
     BINARY_SUBTRACT = stack_op(operator.sub)
-    BINARY_SUBSCR = break_graph_if_unsupported(stack_op(operator.getitem))
+    BINARY_SUBSCR = break_graph_if_unsupported(push=1)(stack_op(operator.getitem))
     BINARY_LSHIFT = stack_op(operator.lshift)
     BINARY_RSHIFT = stack_op(operator.rshift)
     BINARY_AND = stack_op(operator.and_)


Test patch: