Force RNN modules to be inlined (#975)

They call Tensor.set_ internally with Storage, which is no go for AOTAutograd.
Inline into them so that we can graph break.

Fixes pytorch/functorch#586

Test strategy:

```
./benchmarks/torchbench.py --inductor  -dcuda --no-skip -k tts_angular
```

Note that inductor is still failing, but differently, after this PR.

Signed-off-by: Edward Z. Yang <[email protected]>

Signed-off-by: Edward Z. Yang <[email protected]>

Author: ezyang

torchdynamo/allowed_functions.py

@@ -116,15 +116,20 @@ def _allowed_function_ids():
 
     def _is_allowed_module_prefix(obj):
         allowed_modules = ("torch", "math")
-        disallowed_modules = "torch.optim."
+        # torch.nn.modules.rnn is disallowed because these modules internally
+        # flatten their parameters.  This flattening process will call
+        # Tensor.set_ with a Storage, and Storages cannot be traced with
+        # AOTAutograd; so we need to graph-break. To ensure this, we inline
+        # these functions, rather than keep them opaque-ly in the graph.
+        disallowed_modules = ("torch.optim.", "torch.nn.modules.rnn.")
         allowed_modules_dot = tuple([x + "." for x in allowed_modules])
         module = inspect.getmodule(obj)
         if module is None:
             return False
 
         mod_name = module.__name__
 
-        if mod_name.startswith(disallowed_modules):
+        if any(mod_name.startswith(m) for m in disallowed_modules):
             return False
 
         return mod_name in allowed_modules or mod_name.startswith(allowed_modules_dot)