Repository: pytorch
Issue: Make our source attribution debug prints more useful for Compiler Explorer
Labels: triaged,enhancement,oncall: pt2

### 🐛 Describe the bug

Just for fun, I decided to redo an experiment @msaroufim did a year ago, which was how could we integrate PT2 with Compiler Explorer. I had a few objectives:

1. Make @Chillee's "shape viewer"; that is to say, write some PyTorch code, and then get it annotated variable by variable saying what the sizes of all intermediates are (without having to, e.g., insert prints everywhere). Should work with symbolic sizes.
2. Make a nice frontend for https://gist.github.com/ezyang/192c1b5eb57ff95a46d3b50aa46e3193 ; put in some PyTorch code, annotate it with ranks, find valid input shapes for the function.
3. Take various IRs that torch.compile stack can produce and let you introspect it, the same way usual Godbolt works

I cannibalized the existing Python bytecode disassembler to call into torch.compiler https://gist.github.com/ezyang/bb4bbb5b54eebbff1dc8fd7583c5d2bc which was pretty easy. To setup line number correspondences, I have to update processAsm to parse out line numbers and annotate lines of the produced output with them:

```
    override async processAsm(result) {
        const lineRe = /^\s{0,4}(\d+)(.*)/;
        
        const bytecodeLines = result.asm.split('\n');
        
        const bytecodeResult: ParsedAsmResultLine[] = [];
        let lastLineNo: number | undefined;
        let sourceLoc: AsmResultSource | null = null;
    
        for (const line of bytecodeLines) {
            const match = line.match(lineRe);
            
            if (match) {
                const lineno = parseInt(match[1]);
                sourceLoc = {line: lineno, file: null};
                lastLineNo = lineno;
            } else if (line) {
                sourceLoc = {line: lastLineNo, file: null};
            } else {
                sourceLoc = {line: undefined, file: null};
                lastLineNo = undefined;
            }

            bytecodeResult.push({text: line, source: sourceLoc});
        }

        return {asm: bytecodeResult};
    }
```

Our logs are not super well equipped for doing this.

Let's start with (3). To do this, I need two ingredients: I need to print out the IR, and the IR needs to be annotated with line numbers. None of our log outputs accessible by TORCH_LOGS do this.

TORCH_LOGS=graph_code (this is the closest, but the attribution is ... totally useless?)

```
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG] TRACED GRAPH
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]  ===== __compiled_fn_0 =====
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]  <eval_with_key>.0 class GraphModule(torch.nn.Module):
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]     def forward(self, L_args_0_ : torch.Tensor):
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_args_0_ = L_args_0_
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: /Users/ezyang/miniconda3/lib/python3.8/site-packages/torch/_dynamo/external_utils.py:17, code: return fn(*args, **kwargs)
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul = l_args_0_ * 2;  l_args_0_ = None
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: /Users/ezyang/miniconda3/lib/python3.8/site-packages/torch/_dynamo/external_utils.py:17, code: return fn(*args, **kwargs)
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         add = mul + 1;  mul = None
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         return (add,)
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-29 22:57:03,698] torch._dynamo.output_graph.__graph_code: [DEBUG]
``` 

TORCH_LOGS=aot_graphs (same problem as graph_code)

```
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO] TRACED GRAPH
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]  ===== Forward graph 0 =====
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]  <eval_with_key>.4 from /Users/ezyang/miniconda3/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py:477 in wrapped class <lambda>(torch.nn.Module):
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]     def forward(self, arg0_1: f32[2]):
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]         # File: /Users/ezyang/miniconda3/lib/python3.8/site-packages/torch/_dynamo/external_utils.py:17, code: return fn(*args, **kwargs)
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]         mul: f32[2] = torch.ops.aten.mul.Tensor(arg0_1, 2);  arg0_1 = None
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]         
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]         # File: /Users/ezyang/miniconda3/lib/python3.8/site-packages/torch/_dynamo/external_utils.py:17, code: return fn(*args, **kwargs)
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]         add: f32[2] = torch.ops.aten.add.Tensor(mul, 1);  mul = None
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]         return (add,)
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO]         
[2023-07-29 23:16:30,543] torch._functorch.aot_autograd.__aot_graphs: [INFO] 
```

TORCH_LOGS=output_code

```
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] extern "C" void kernel(const float* in_ptr0,
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]                        float* out_ptr0)
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] {
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     {
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]         #pragma GCC ivdep
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]         for(long i0=static_cast<long>(0L); i0<static_cast<long>(2L); i0+=static_cast<long>(1L))
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]         {
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]             auto tmp0 = in_ptr0[static_cast<long>(i0)];
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]             auto tmp1 = static_cast<float>(2.0);
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]             auto tmp2 = decltype(tmp0)(tmp0 * tmp1);
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]             auto tmp3 = static_cast<float>(1.0);
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]             auto tmp4 = tmp2 + tmp3;
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]             out_ptr0[static_cast<long>(i0)] = tmp4;
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]         }
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     }
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] }
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] ''')
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] 
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] 
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] async_compile.wait(globals())
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] del async_compile
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] 
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG] def call(args):
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     arg0_1, = args
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     args.clear()
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     assert_size_stride(arg0_1, (2, ), (1, ))
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     buf0 = empty_strided((2, ), (1, ), device='cpu', dtype=torch.float32)
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     cpp_fused_add_mul_0(c_void_p(arg0_1.data_ptr()), c_void_p(buf0.data_ptr()))
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     del arg0_1
[2023-07-29 22:59:22,230] torch._inductor.graph.__output_code: [DEBUG]     return (buf0, )
```

We do have a print for tensor sizes, there's no way to directly get at it as an artifact, but anyway it looks like this:

```
[2023-07-29 23:17:55,235] torch._dynamo.output_graph.__graph_sizes: [DEBUG] TRACED GRAPH TENSOR SIZES
[2023-07-29 23:17:55,235] torch._dynamo.output_graph.__graph_sizes: [DEBUG] ===== __compiled_fn_0 =====
[2023-07-29 23:17:55,235] torch._dynamo.output_graph.__graph_sizes: [DEBUG] l_args_0_: (2,)
[2023-07-29 23:17:55,235] torch._dynamo.output_graph.__graph_sizes: [DEBUG] mul: (2,)
[2023-07-29 23:17:55,235] torch._dynamo.output_graph.__graph_sizes: [DEBUG] add: (2,)
```

So what are the problems?

1. Source code attribution seems to not work if you compile()/exec() some source code. I can manually work around this by using importloader to load the module to execute, so this isn't a hard blocker
2. Many of the IR outputs don't print line numbers. They should, so that I can setup the correspondences easily (inductor code output, tensor sizes)

cc @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @wconstab @williamwen42 @voznesenskym 

### Versions

main