[rfc] add debug mode to print meta in fx graphs

ghstack-source-id: 48e5628e3f6cec39a6f997c5ac2a7731fd3188f1
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165874

torch/fx/graph.py

@@ -454,6 +454,7 @@ class CodeGen:
         include_device = include_device or (
             os.environ.get("FX_GRAPH_SHOW_DEVICE", "0") == "1"
         )
+        include_meta = os.environ.get("FX_GRAPH_SHOW_META", "0") == "1"
 
         def add_global(name_hint: str, obj: Any):
             """Add an obj to be tracked as a global.
@@ -689,6 +690,13 @@ class CodeGen:
                     maybe_comment += f"  # {desc}"
                 # output is handled specially
 
+            if include_meta and hasattr(node, "meta") and node.meta:
+                # Print the node.meta dict as a comment above the node
+                import pprint
+                meta_str = pprint.pformat(node.meta, width=80, compact=True)
+                for line in meta_str.splitlines():
+                    body.append(f"# meta: {line}\n")
+
             if node.op == "placeholder":
                 assert isinstance(node.target, str)
                 maybe_default_arg = (