Output runnable backwards _inductor graph `torch.compile`

Jerome_Ku: <h3><a name="the-feature-motivation-and-pitch-1" class="anchor" href="#the-feature-motivation-and-pitch-1"></a><img src="https://discuss.pytorch.org/images/emoji/apple/rocket.png?v=12" title=":rocket:" class="emoji" alt=":rocket:" loading="lazy" width="20" height="20"> The feature, motivation and pitch</h3>
<p>How can I get a runnable copy of the backwards graph when running <code>torch.compile</code>?</p>
<p>For example:</p>
<pre data-code-wrap="python"><code class="lang-python">@torch.compile
def f(x):
    out = x.sin() + x.cos()
    return out

x = torch.ones(2, requires_grad=True).cuda()
out = f(x)
</code></pre>
<p>If I run this with <code>TORCH_LOGS=all</code>, I can see that the backwards graph is traced:</p>
<pre><code class="lang-auto">===== Backward graph 0 =====
 &lt;eval_with_key&gt;.36 class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[2]", tangents_1: "f32[2]"):
        # File: /notebooks/clean-repos/triton-autodiff/test_compile.py:32 in f, code: tmp1 = x.sin() + x.cos()
        sin: "f32[2]" = torch.ops.aten.sin.default(primals_1)
        cos: "f32[2]" = torch.ops.aten.cos.default(primals_1);  primals_1 = None
        neg: "f32[2]" = torch.ops.aten.neg.default(sin);  sin = None
        mul: "f32[2]" = torch.ops.aten.mul.Tensor(tangents_1, neg);  neg = None
        mul_1: "f32[2]" = torch.ops.aten.mul.Tensor(tangents_1, cos);  tangents_1 = cos = None
        
        # File: /notebooks/clean-repos/triton-autodiff/test_compile.py:32 in f, code: tmp1 = x.sin() + x.cos()
        add_1: "f32[2]" = torch.ops.aten.add.Tensor(mul, mul_1);  mul = mul_1 = None
        return [add_1]
</code></pre>
<p>However, if I look at the <code>inductor</code> cache, I see only the forward kernel (<code>output.py</code>):</p>
<pre data-code-wrap="python"><code class="lang-python">
from ctypes import c_void_p, c_long
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align

from torch import device, empty, empty_strided
from torch._inductor.codecache import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
from torch._inductor.codegen.multi_kernel import MultiKernelCall

aten = torch.ops.aten
inductor_ops = torch.ops.inductor
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
alloc_from_pool = torch.ops.inductor._alloc_from_pool
reinterpret_tensor = torch.ops.inductor._reinterpret_tensor
async_compile = AsyncCompile()


# kernel path: torch_compile/cache/rv/crviokfbw7lr4fi4yf2kqhtyjostdohc32ulg7vlfbyprvjinxaa.py
# Source Nodes: [cos, sin, tmp1], Original ATen: [aten.add, aten.cos, aten.sin]
# cos =&gt; cos
# sin =&gt; sin
# tmp1 =&gt; add
triton_poi_fused_add_cos_sin_0 = async_compile.triton('triton_', '''
import triton
import triton.language as tl
from torch._inductor.ir import ReductionHint
from torch._inductor.ir import TileHint
from torch._inductor.triton_heuristics import AutotuneHint, pointwise
from torch._inductor.utils import instance_descriptor
from torch._inductor import triton_helpers

@pointwise(
    size_hints=[2], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [instance_descriptor(divisible_by_16=(0, 1), equal_to_1=(), ids_of_folded_args=(), divisible_by_8=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cos_sin_0', 'mutated_arg_names': [], 'no_x_dim': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 2
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex &lt; xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp1 = tl.sin(tmp0)
    tmp2 = tl.cos(tmp0)
    tmp3 = tmp1 + tmp2
    tl.store(out_ptr0 + (x0), tmp3, xmask)
''')
...
</code></pre>
<p>How can I get the <code>codegen'ed</code> equivalents for the backwards graph, and the <code>torch.autograd.Function</code> that calls these forward / backward kernels?</p>