Very first node in FX graph has misattributed stack

[ezyang]

Steps to reproduce:

1. Build PyTorch at Symbolic shapes pytorch#84246 (at time of writing, I was on ce220330e5d69899ee89ea8dd533ae81e08e0667 ); this step may not be necessary but I did not try
2. Make log.info print stack traces for the traced graph with this patch

diff --git a/torchdynamo/output_graph.py b/torchdynamo/output_graph.py
index db84c3e..54d26a3 100644
--- a/torchdynamo/output_graph.py
+++ b/torchdynamo/output_graph.py
@@ -373,7 +373,7 @@ class OutputGraph(fx.Tracer):
             # the call to tabulate can cause a lot of memory to be allocated
             if config.log_level <= logging.INFO:
                 log.info(
-                    f"TRACED GRAPH\n {name} {gm.forward.__code__.co_filename} {format_graph_tabular(gm.graph)}\n"
+                    f"TRACED GRAPH\n {name} {gm.forward.__code__.co_filename} {gm.print_readable()}\n"
                 )
         except ImportError:
             log.warning(

(btw, we should consider making this the default)

3. Inspect the printed graph from python benchmarks/torchbench.py --only BERT_pytorch --accuracy-aot-nop:

    def forward(self, inputs_0_ : torch.Tensor, inputs_1_ : torch.Tensor):                                      
                                                                                                                
        # Module stack: {'mod': 'BERT'}, File: benchmarks/torchbench.py:336, code: return mod(*inputs)          
        gt = inputs_0_ > 0                                                                                      
                                                                                                                
        # Module stack: {'mod': 'BERT'}, File: /raid/ezyang/benchmark/torchbenchmark/models/BERT_pytorch/bert_pytorch/model/bert.py:40, code: mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)                  
        unsqueeze = gt.unsqueeze(1);  gt = None                                                                 
        repeat = unsqueeze.repeat(1, 128, 1);  unsqueeze = None                                                 
        unsqueeze_1 = repeat.unsqueeze(1);  repeat = None

You can see that the inputs_0_ > 0 operation has a useless stack trace, and the next set of lines (unsqueeze, repeat, unsqueeze) have the correct stack trace (luckily for me, this is the right stack trace for gt too).

This suggests that dynamo doesn't have the correct stack trace setup at the very beginning of a graph capture.

cc @SherlockNoMad

[williamwen42]

Minimum repro:

import torch
import torchdynamo

def fn(x):
    y = x > 0
    z = y + 1
    return z
    
torchdynamo.config.log_level = torchdynamo.logging.CODE

x = torch.randn(10)

@torchdynamo.optimize()
def g(x):
    return fn(x)

g(x)

produces the following graph:

class GraphModule(torch.nn.Module):
    def forward(self, x : torch.Tensor):
        
        # Module stack: {}, File: test_stack_trace.py:17, code: return fn(x)
        gt = x > 0;  x = None
        
        # Module stack: {}, File: test_stack_trace.py:8, code: z = y + 1
        add = gt + 1;  gt = None
        return (add,)