Repository: pytorch
Issue: Compiled autograd: log which saved activations came from which fw graph regions
Labels: triaged,oncall: pt2,module: dynamo,module: compiled autograd

One piece of compiled autograd debugging that I think might aid in debugging would be that:
* in the "non-compiled autograd flow", when you inspect a given backward graph and I see some input that is not named `tangents*`, I know that it is a saved activation
* I can inspect the corresponding forward graph to see exactly how it was computed, and what user code it was computed from (through the stack trace info)
* This info is lost in the compiled autograd graph: in the compiled autograd graph, you see a big list of inputs, but you do not know where they came from.

One thing that might be nice would be something like: when compiled autograd pulls a bunch of saved activations from a given node on the autograd tape, it checks to see if that node is a `CompiledBackwardNode`. If it is, also log for that particular input:
(1) the aot_id of this backward node (so you can correlate it with the forward - we can store the aot_id directly on the node: https://github.com/pytorch/pytorch/blob/main/torch/_functorch/_aot_autograd/runtime_wrappers.py#L1477)
(2) the index of this particular activation from that backward node (so we can map it to, e.g. "output 7 of forward graph 5")

example code:
```
import torch

@torch.compile
def f(x):
    tmp1 = x.sin()
    print("graph break!")
    return tmp1.sin()

x = torch.randn(4, requires_grad=True)
out = f(x)
with torch._dynamo.compiled_autograd.enable(torch.compile):
    out.sum().backward()
```

If I run the above with `TORCH_LOGS="compiled_autograd,aot"`, I get the below graphs. I happen to know that in the compiled autograd graph, `inputs[1]` and `inputs[2]` correspond to the saved activation in forward graphs 1 and 2 (`primals_1`, `mul` nodes), but it would be nice to be able to tell that strictly from the logs.
```
 ===== Forward graph 0 =====
 /home/hirsheybar/local/a/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[4][1]cpu"):
         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:6 in f, code: tmp1 = x.sin()
        sin: "f32[4][1]cpu" = torch.ops.aten.sin.default(primals_1)
        return (sin, primals_1)


INFO: TRACED GRAPH
 ===== Backward graph 0 =====
 <eval_with_key>.1 class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[4][1]cpu", tangents_1: "f32[4][1]cpu"):
         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:6 in f, code: tmp1 = x.sin()
        cos: "f32[4][1]cpu" = torch.ops.aten.cos.default(primals_1);  primals_1 = None
        mul: "f32[4][1]cpu" = torch.ops.aten.mul.Tensor(tangents_1, cos);  tangents_1 = cos = None
        return (mul,)

 ===== Forward graph 1 =====
 /home/hirsheybar/local/a/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[4][1]cpu"):
         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:8 in torch_dynamo_resume_in_f_at_7, code: return tmp1.sin()
        sin: "f32[4][1]cpu" = torch.ops.aten.sin.default(primals_1)
        return (sin, primals_1)


INFO: TRACED GRAPH
 ===== Backward graph 1 =====
 <eval_with_key>.5 class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[4][1]cpu", tangents_1: "f32[4][1]cpu"):
         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:8 in torch_dynamo_resume_in_f_at_7, code: return tmp1.sin()
        cos: "f32[4][1]cpu" = torch.ops.aten.cos.default(primals_1);  primals_1 = None
        mul: "f32[4][1]cpu" = torch.ops.aten.mul.Tensor(tangents_1, cos);  tangents_1 = cos = None
        return (mul,)


INFO: TRACED GRAPH
 ===== Compiled autograd graph =====
 <eval_with_key>.8 class CompiledAutograd(torch.nn.Module):
    def forward(self, inputs, sizes, hooks):
        # No stacktrace found for following nodes
        getitem: "f32[]" = inputs[0]
        getitem_1: "f32[4]" = inputs[1]
        getitem_2: "f32[4]" = inputs[2];  inputs = None
        expand: "f32[4]" = torch.ops.aten.expand.default(getitem, [4]);  getitem = None
        clone: "f32[4]" = torch.ops.aten.clone.default(expand, memory_format = torch.contiguous_format);  expand = None
        cos: "f32[4]" = torch.ops.aten.cos.default(getitem_1);  getitem_1 = None
        mul: "f32[4]" = torch.ops.aten.mul.Tensor(clone, cos);  clone = cos = None
        cos_1: "f32[4]" = torch.ops.aten.cos.default(getitem_2)
        mul_1: "f32[4]" = torch.ops.aten.mul.Tensor(mul, cos_1);  mul = cos_1 = None
        accumulate_grad_ = torch.ops.inductor.accumulate_grad_.default(getitem_2, mul_1);  getitem_2 = mul_1 = accumulate_grad_ = None
        _exec_final_callbacks_stub = torch__dynamo_external_utils__exec_final_callbacks_stub();  _exec_final_callbacks_stub = None
        return []


INFO: TRACED GRAPH
 ===== Forward graph 2 =====
 /home/hirsheybar/local/a/pytorch/torch/fx/_lazy_graph_module.py class <lambda>(torch.nn.Module):
    def forward(self, arg0_1: "f32[][]cpu", arg1_1: "f32[4][1]cpu", arg2_1: "f32[4][1]cpu"):
        # No stacktrace found for following nodes
        expand: "f32[4][0]cpu" = torch.ops.aten.expand.default(arg0_1, [4]);  arg0_1 = None
        clone: "f32[4][1]cpu" = torch.ops.aten.clone.default(expand, memory_format = torch.contiguous_format);  expand = None
        cos: "f32[4][1]cpu" = torch.ops.aten.cos.default(arg1_1);  arg1_1 = None
        mul: "f32[4][1]cpu" = torch.ops.aten.mul.Tensor(clone, cos);  clone = cos = None
        cos_1: "f32[4][1]cpu" = torch.ops.aten.cos.default(arg2_1);  arg2_1 = None
        mul_1: "f32[4][1]cpu" = torch.ops.aten.mul.Tensor(mul, cos_1);  mul = cos_1 = None
        clone_1: "f32[4][1]cpu" = torch.ops.aten.clone.default(mul_1);  mul_1 = None
        return (clone_1,)
```

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames