Update how Dynamo decides to graph break on an OpOverloadPacket (#112200)

Summary:
Previously, under config.only_allow_pt2_compliant_ops, Dynamo graph
breaks when it see an OpOverloadPacket where any overloads are not
PT2 compliant. This is potentially brittle: if someone (unlikely) adds
a new overload for a custom operator, then this would cause a
previously non-graph-breaking call to the OpOverloadPacket to graph break.

In this PR:
- When Dynamo is about to write a call to an operator to the FX graph,
we check if it is PT2 compliant.
- For OpOverload, we check to see if the tag is on it
- For OpOverloadPacket, we do overload resolution and check to see if
  the tag is on the OpOverload that it resolves to.

X-link: pytorch/pytorch#112200
Approved by: https://github.com/bdhirsh

Reviewed By: ZainRizvi

Differential Revision: D50873052

Pulled By: zou3519

fbshipit-source-id: c95b9797185f4bfc9edd2eb2e48f73b8dd56154f

Author: zou3519

userbenchmark/dynamo/dynamobench/_dynamo/utils.py

@@ -1387,6 +1387,21 @@ def extract_fake_example_value(node, required=True):
         return None
 
 
+def ensure_graph_fake(e, tx):
+    assert maybe_get_fake_mode(e) is tx.fake_mode
+    return e
+
+
+def get_fake_values_from_nodes(tx, nodes):
+    def visit(n: torch.fx.Node):
+        return n.meta["example_value"]
+
+    args_kwargs = torch.fx.node.map_arg(nodes, visit)
+    return tree_map_only(
+        torch.Tensor, functools.partial(ensure_graph_fake, tx=tx), args_kwargs
+    )
+
+
 def get_fake_value(node, tx, allow_non_graph_fake=False):
     """
     Run the computation represented by `node` using fake tensors and return the result.
@@ -1410,16 +1425,7 @@ def get_fake_value(node, tx, allow_non_graph_fake=False):
     if "example_value" in node.meta and is_fake(node.meta["example_value"]):
         return node.meta["example_value"]
 
-    def ensure_graph_fake(e):
-        assert maybe_get_fake_mode(e) is tx.fake_mode
-        return e
-
-    def visit(n: torch.fx.Node):
-        return n.meta["example_value"]
-
-    args, kwargs = torch.fx.node.map_arg((node.args, node.kwargs), visit)
-    args = tree_map_only(torch.Tensor, ensure_graph_fake, args)
-    kwargs = tree_map_only(torch.Tensor, ensure_graph_fake, kwargs)
+    args, kwargs = get_fake_values_from_nodes(tx, (node.args, node.kwargs))
 
     nnmodule = None
     if op == "call_method" and len(args) > 0 and isinstance(args[0], torch.nn.Module):
@@ -1483,7 +1489,9 @@ def visit(n: torch.fx.Node):
         raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
 
     if not allow_non_graph_fake:
-        _ = tree_map_only(torch.Tensor, ensure_graph_fake, ret_val)
+        _ = tree_map_only(
+            torch.Tensor, functools.partial(ensure_graph_fake, tx=tx), ret_val
+        )
     return ret_val