[export] Make dict_keys_getitem tracable. (pytorch#166776)

Summary:
dict_keys_getitem can show up in the bytecode but it's using dict.keys() which is not fx tracable.

fx.wrap should make it as a standalone function in the graph to be invoked later with real inputs.

Test Plan:
pytest test/export/test_experimental.py

Pull Request resolved: pytorch#166776
Approved by: https://github.com/jamesjwu
ghstack dependencies: pytorch#166775

Author: zhxchen17

test/export/test_experimental.py

@@ -541,6 +541,44 @@ def make_inputs(b: int):
         self.assertIsNotNone(gm.meta["tracing_context"].fake_mode)
         self.assertEqual(len(gm.meta["tracing_context"].tensor_to_context), 1)
 
+    def test_dynamo_graph_capture_dict_keys_getitem(self):
+        class Module(torch.nn.Module):
+            def forward(self, x):
+                return x * 2
+
+        foo = Module()
+
+        class BlockMask:
+            def __init__(self, d):
+                self.d = d
+
+        block_mask = BlockMask(torch.randn(4))
+
+        def pre_hook_function(m, input):
+            block_mask.d = input[0] + 1
+            return input  # Return a tuple of modified inputs
+
+        foo.register_forward_pre_hook(pre_hook_function)
+
+        def make_inputs():
+            return (torch.randn(4),)
+
+        trace_inputs = make_inputs()
+        gm = dynamo_graph_capture_for_export(foo)(*trace_inputs)
+        test_inputs = make_inputs()
+        self.assertExpectedInline(
+            gm.code.strip("\r\n "),
+            """\
+def forward(self, args_0):
+    _tree_leaf_0, _tree_leaf_1, = pytree.tree_leaves((self, args_0,))
+    L_args_0_ , = self._in_shuffle_graph(_tree_leaf_0, _tree_leaf_1)
+    l_args_0_ = L_args_0_
+    add = l_args_0_ + 1
+    mul = l_args_0_ * 2;  l_args_0_ = None
+    return pytree.tree_unflatten(self._out_shuffle_graph(_tree_leaf_0, _tree_leaf_1, mul, add), self._out_spec)""",
+        )
+        self.assertEqual(gm(*test_inputs), foo(*test_inputs))
+
     @unittest.skipIf(not TEST_CUDA, "CUDA not available")
     def test_dynamo_graph_capture_fx_graph_annotate_overlap_pass(self):
         class DummyOp(torch.autograd.Function):

torch/_dynamo/convert_frame.py

@@ -1000,7 +1000,10 @@ def get_traced_fn(mod: Any) -> tuple[FunctionType, Optional[object]]:
     import inspect
 
     if isinstance(mod, torch.nn.Module):
-        mod = mod.forward
+        if len(mod._forward_pre_hooks) == 0 and len(mod._forward_hooks) == 0:
+            mod = mod.forward
+        else:
+            mod = mod.__call__
     if hasattr(mod, "__self__"):
         # pyrefly: ignore [missing-attribute]
         return mod.__func__, mod.__self__

torch/_dynamo/functional_export.py

@@ -526,6 +526,8 @@ def backend_dummy(*example_inputs):
     in_shuffle_graph = make_fx(
         InShuffle(), tracing_mode="symbolic", proxy_module_inputs=True
     )(*flat_real_args)
+    in_shuffle_graph.graph.eliminate_dead_code()
+    in_shuffle_graph.recompile()
 
     output_node = next(iter(reversed(backend_input.graph_module.graph.nodes)))
 
@@ -573,6 +575,8 @@ def backend_dummy(*example_inputs):
     out_shuffle_graph = make_fx(
         out_shuffle, tracing_mode="symbolic", proxy_module_inputs=True
     )(*flat_out_shuffle_args)
+    out_shuffle_graph.graph.eliminate_dead_code()
+    out_shuffle_graph.recompile()
 
     assert out_shuffle.out_spec is not None
     return PyTreeifyOutput(

torch/_dynamo/utils.py

@@ -2707,6 +2707,7 @@ def to_subclass(t: Any, cls: type) -> Any:
 dict_getitem = dict.__getitem__
 
 
+@torch.fx.wrap
 def dict_keys_getitem(d: dict[Any, Any], n: int) -> Any:
     # Call dict(d) to prevent calling overridden __iter__/keys
     dict_class = dict