Torch.compile Deep Dive II — Inductor Codegen and Buffer Lifetimes, End to End

All permalinks in this post are pinned to pytorch/pytorch commit e2f06e4 (April 2026). Line numbers have been verified against that revision.

🧭 Part 2 of a torch.compile series Part 1 — Torch.compile 101: From Python Function to Triton Kernel walked through the high-level pipeline: TorchDynamo’s bytecode-level graph capture, AOTAutograd’s joint forward/backward graph, and a tour of Inductor’s main phases. If you haven’t read it, the short version: by the time Inductor enters the picture, you have a flat FX graph of aten.* ops on fake tensors, and Inductor’s job is to turn that into a Python wrapper that allocates buffers and launches Triton kernels. This post picks up there and goes deep.

Hands-on companions in torch-compile-tutorial:

• Lesson 4 (04_inductor_codegen.{py,md}) — walks one tiny function through every stage with the actual generated wrapper and kernel source side-by-side.
• deep_dive_ii_inductor_internals/ — drop-in inductor_hooks.py that prints one line per Inductor stage as compilation runs (call_function dispatch → realize → register_buffer → scheduler _compute_attrs → compute_last_usage → wrapper codegen_free → Triton codegen_kernel). Maps 1:1 to blog stages 1–6.

Introduction

Most write-ups on Inductor walk you through the stages — Dynamo, AOTAutograd, IR lowering, scheduler, codegen — and treat dependency tracking as one item on the list. This post does the opposite. It threads a single question through the entire pipeline: when does Inductor decide a buffer is no longer needed, and where does that decision come from? That single question turns out to surface almost everything else worth knowing — how IR is built, how read_writes is computed, how fusion respects dependencies, how the wrapper emits del bufN, how kernel signatures get assembled.

Buffer lifetimes are also the most directly observable thing about a compiled model. Open any output_code.py and you see a stream of bufN = empty_strided_cuda(...), kernel calls, and del bufN lines. Every one of those lines was produced by a layer in the pipeline reading off the layer below it. Pin down how that read-off works at each stage and you can reverse-engineer any piece of generated code back to its source FX node.

Two design ideas anchor the rest of the post:

1. Closure inlining at lowering time. Most pointwise / reduction ops never get a bufN name. Their compute survives as a captured inner_fn inside the consumer’s closure. By the time the scheduler runs, what was a long chain of ATen ops is one expression. This is what makes “fusion” in Inductor mostly free — the work has already been done.
2. Realization is the trigger for everything else. When (and only when) something forces a TensorBox to realize, a ComputedBuffer is born, register_buffer fires, and the scheduler / wrapper / kernel codegen suddenly have something to look at. So tracking when realize fires tells you exactly which compute lives in registers vs. HBM.

Concretely, we’ll trace one tiny model through six stages, ending with the actual generated wrapper and Triton kernel. All permalinks are pinned to a known pytorch revision so the line numbers stay stable.

The function

def fn(a, b, c):
    return ((a + b) * c.relu()).sum(dim=-1)

Three pointwise ops (add, relu, mul) + one reduction (sum), all on tensors of shape (64, 128). Three graph inputs. One graph output.

Stage 0 — From torch.compile to Inductor’s front door

When you call torch.compile(fn)(...), control flows: Dynamo traces Python bytecode into an FX graph, AOTAutograd runs functionalization and decompositions and splits forward / backward, then hands each half to a backend. With the default backend, that backend is Inductor.

Inductor’s entry point is fx_codegen_and_compile, which receives a torch.fx.GraphModule and example inputs:

def fx_codegen_and_compile(
    gm: GraphModule,
    example_inputs: Sequence[InputType],
    inputs_to_check: Sequence[int],
    compile_region_name: str | None = None,
    **graph_kwargs: Unpack[_CompileFxKwargs],
) -> OutputCode:

The FX graph at this boundary already has all autograd machinery resolved: nodes are placeholder, call_function (typically with aten.* or prims.* targets), and output. There are no torch.nn.Module calls, no autograd tape, no Python control flow — just a flat dataflow graph of ATen operations on fake tensors.

What Inductor does next is, broadly, two passes: an FX-walk that builds Inductor’s own IR (the GraphLowering pass), then a backend pass that schedules, fuses, and emits Python + Triton.

Stage 1 — GraphLowering: FX → Inductor IR

GraphLowering is a torch.fx.Interpreter, so it walks the FX graph node-by-node by graph.run, dispatching on node opcode. There are three different cases:

1. placeholder — graph inputs become InputBuffers

2. output — graph outputs are pinned in GraphLowering.output records which buffers the compiled function will return. These names get added to V.graph.get_output_names() and become the seed for the scheduler’s last-use analysis: outputs are live to the end, so they never get freed inside call(...).

3. call_function — ATen ops become Inductor IR via lowerings, where we need to dive deeper. Our example code was traced into following fx graph:

arg0_1 = placeholder()          # a
arg1_1 = placeholder()          # b
arg2_1 = placeholder()          # c
add_1  = aten.add.Tensor(arg0_1, arg1_1)
relu_1 = aten.relu.default(arg2_1)
mul_1  = aten.mul.Tensor(add_1, relu_1)
sum_1  = aten.sum.dim_IntList(mul_1, [-1])
output (sum_1,)

Four call_function nodes. For each, GraphLowering.call_function looks up the FX target in the global lowerings registry:

lowerings: dict[Callable[..., Any] | str, Callable[..., Any]] = {}

Hundreds of entries, registered via the @register_lowering(...) decorator. Take the first node, add_1 = aten.add.Tensor(arg0_1, arg1_1). The registry entry was created at module-load time by:

add = register_pointwise(aten.add, allow_alpha=True)

Inside register_pointwise:

def register_pointwise(aten_fn, name=None, ...):
    name = name or aten_fn.__name__               # "add"
    fn = ops_wrapper(name)                        # closure tied to that name
    ...
    return register_lowering(aten_fn, broadcast=broadcast, ...)(
        make_pointwise(fn, ...)
    )

make_pointwise(fn, ...) is a higher-order factory. It returns an inner(*inputs) callable that, when invoked by the dispatcher, builds the actual IR node. make_pointwise itself is op-agnostic — the op identity ("add" vs "mul" vs "relu") is captured entirely in the fn closure passed in.

So when GraphLowering.call_function reaches add_1:

1. Look up lowerings[aten.add.Tensor] → returns the inner callable produced by make_pointwise(ops_wrapper("add")).

2. Call inner(tb_arg0_1, tb_arg1_1). Inside, make_pointwise’s inner_fn is constructed:

   loaders = [tb_arg0_1.make_loader(), tb_arg1_1.make_loader()]
   def inner_fn(index):
       inputs_loaded = [load(index) for load in loaders]
       out = fn(*inputs_loaded)              # fn = ops_wrapper("add") → ops.add(...)
       return out
   return Pointwise.create(device=..., dtype=..., inner_fn=inner_fn, ranges=...)
   

3. Pointwise.create wraps the new Pointwise in a TensorBox: return TensorBox.create(Pointwise(...)).
4. The TensorBox is returned from call_function. FX’s Interpreter stores it in self.env[add_1] for downstream nodes to pick up.

The “result” of aten.add.Tensor is purely a recipe — a closure over an output index that, when called with [i, j, ...], would compute a[i,j] + b[i,j].

💡 Side note: not every op goes through make_pointwise Pointwise + reduction ops dominate by count, but a handful of FX targets dispatch to entirely different IR types:

• Matmul-style ops (aten.mm, aten.bmm, aten.convolution, aten._scaled_dot_product_flash_attention, flex_attention, …) lower to a TritonTemplateBuffer (or a sibling like CuteDSLTemplateBuffer). The body is a Jinja template, not a Python inner_fn, so the scheduler can’t symbolically interpret it. These call register_buffer immediately at construction.
• External library calls (cuBLAS, cuDNN, mkldnn, NCCL collectives, custom ops) lower to ExternKernel subclasses, which also register at construction. The kernel call site emits a torch.ops.<lib>.<name>(...) line in the wrapper.
• Concat / index_put / scatter lower to specialized kernels (ConcatKernel, mutation-aware paths) that bypass the closure form because they have multi-output or in-place semantics.

In our worked example none of these fire — every call_function is pointwise or reduction, so the closure-based make_pointwise / make_reduction path handles all four.

The same chain repeats for relu_1, mul_1, and the body of sum_1 (the Reduction’s epilogue is also a Pointwise). The entire call_function pass produces four nested TensorBoxes whose inner_fns recursively close over each other — but no buffer has been registered yet. Realization happens lazily

When does realization actually fire?

So far we’ve said the closure stays dormant “until something realizes it.” Three things to pin down before that handwave is useful: what realize does, what register_buffer does, and when the closure actually runs.

What realize() does

The realization API is TensorBox.realize:

def realize(self) -> str | None:
    ...
    self.data = ComputedBuffer(
        name=None,
        layout=FlexibleLayout(device=device, dtype=..., size=..., is_pinned=False),
        data=self.data,                              # the Pointwise / Reduction chain
    )
    self.data.name = V.graph.register_buffer(self.data)
    V.graph.register_operation(self.data)
    ...
    return self.data.name

Two things happen, in order. First, the existing Pointwise/Reduction chain (already living inside self.data) is wrapped in a ComputedBuffer that carries a FlexibleLayout. The inner_fn is not changed — it still closes over the same upstream Pointwisees. The ComputedBuffer is just an outer wrapper that adds shape/stride information so the buffer can be allocated. Second, register_buffer is called.

💡 What is V.graph? V is Inductor’s thread-local virtualized state holder (from torch._inductor.virtualized import V). Three slots get used heavily in this post:

• V.graph — the current GraphLowering instance. compile_fx enters its main region with with V.set_graph_handler(graph_lowering): ..., and from that point on every V.graph.register_buffer(...), V.graph.buffers, V.graph.get_output_names() resolves to that one instance.
• V.ops — the currently-installed virtualized ops handler. ops.add(a, b) inside an inner_fn dispatches through V.ops to whatever is active: a tracking handler during extract_read_writes, TritonOverrides during Triton codegen, CppOverrides during C++ codegen. Same inner_fn, different output.
• V.kernel — the active Kernel object during codegen (carries V.kernel.cse, V.kernel.args, V.kernel.body).

All three are dynamically scoped via with-block context managers, so they nest cleanly. There’s no module-level graph variable anywhere in Inductor — every reference to “the current graph” goes through V.graph, which keeps the code reentrant and lets test code install fake handlers.

What register_buffer does

Short and load-bearing — graph.py:1070:

def register_buffer(self, buffer, *, set_name=False) -> str:
    name = self.qualify_name(f"buf{len(self.buffers)}")   # monotonic: buf0, buf1, ...
    self.buffers.append(buffer)
    self.name_to_buffer[name] = buffer
    ...
    return name

It mints a fresh bufN name by counting self.buffers, appends the buffer to the list, and adds a name→buffer entry to the dict. That’s it. No allocation happens. No memory is touched. No kernel is emitted. What changes is purely that V.graph.buffers and V.graph.name_to_buffer now know this buffer exists.

But that’s exactly the visibility every downstream stage requires:

• The scheduler walks V.graph.buffers (in Scheduler._init) to wrap each entry in a SchedulerNode.
• The wrapper codegen looks up name_to_buffer[name] to emit bufN = empty_strided_cuda(...) allocation lines.
• The kernel codegen emits the buffer’s bufN string into pointer-arg positions.

Before register_buffer, none of that can happen — nobody else in the pipeline knows the buffer is there. The ComputedBuffer exists in memory only as an attribute on a TensorBox that’s about to be GC’d from FX’s env. After register_buffer, it’s anchored in V.graph and visible everywhere.

When the closure is actually called

A subtle point: building the Pointwise chain at lowering time does not invoke any inner_fn. The closure is constructed but dormant. It survives only because each consumer’s inner_fn closes over the producer’s inner_fn (via make_loader). At lowering time, no ops.load or ops.add ever runs; we’re just recording compose rules.

The closure only fires later, during two specific passes that install a virtualized ops handler and call the terminal inner_fn(index):

1. Stage 2 — SchedulerNode._compute_attrs symbolically interprets the body to extract read_writes. The handler’s ops.load(name, idx) records a read; its ops.store(name, idx) records a write. Calling Reduction_S.inner_fn(idx) cascades into every captured upstream inner_fn, and the recorded reads/writes settle on the leaves — graph inputs and other realized buffers.
2. Stage 6 — TritonKernel.codegen_body emits the kernel string. The handler is TritonOverrides: ops.load(...) returns "tl.load(in_ptr0 + idx, mask)", ops.add(a, b) returns "tmp = a + b", etc. Same closure tree, different handler, different output.

These are the only two times the closure runs. The first one is what gives us read_writes (and therefore everything in stages 3–5); the second is what produces the kernel source.

When realize() itself is invoked

1. From GraphLowering.run_node (the FX walker), at three trigger points in run_node:

• The FX node has multiple users and one of them is in needs_realized_inputs (ops that can’t take a Pointwise — aten.mm, conv, scatter, …) → calls realize_hint() which realizes only if the closure has more than one nontrivial read.
• The FX node carries meta["inductor_realize_to_strides"] → forces realize() to pin a stride layout.
• The node is a graph output → realizes (subject to config.delay_realize_cheap_outputs).

2. From specific lowerings that can’t operate on Pointwise form. About 20 sites in lowering.py: nonzero, unique, bincount, searchsorted, index_put, scatter/gather variants, etc. Anything that depends on data values (not just indices) or has multi-output semantics needs a real buffer.

3. From mutation handling. When a downstream op mutates an input, GraphLowering.register_users_of walks all users of the about-to-be-mutated buffer and forces each to realize first — otherwise the closure-inlined readers would see post-mutation data.

Plus an eager path for IR types that are always materialized: TemplateBuffer (matmul-style kernels), ExternKernel subclasses (cuBLAS, cuDNN, scatter, custom ops), ConcatKernel — these all call register_buffer directly inside __init__ because they have no inner_fn form to inline.

So there’s a clean dichotomy: if you can be expressed as a closure, you stay unrealized as long as possible; if you can’t, you’re a buffer immediately. Everything in between is governed by the heuristics above. If you set a breakpoint at register_buffer and step through any compile, the call sites you hit are exactly this list.

Walk through 01: realization in the running example

Apply this to ((a + b) * c.relu()).sum(dim=-1):

FX node	TensorBox built	Realize?	Why
add_1	tb_A = TensorBox(Pointwise_A)	no	add_1.users = {mul_1}, single user, mul is pointwise (not in needs_realized_inputs)
relu_1	tb_B = TensorBox(Pointwise_B)	no	same — single pointwise user
mul_1	tb_M = TensorBox(Pointwise_M) whose inner_fn closes over both Pointwise_A.inner_fn and Pointwise_B.inner_fn (via make_loader)	no	single user sum_1, sum is a reduction (not in needs_realized_inputs)
sum_1	tb_S = TensorBox(Reduction_S) whose inner_fn closes over Pointwise_M.inner_fn	no — yet	single user is the graph output
output	GraphLowering.output runs	yes	line 1600 calls ir.ExternKernel.realize_input(x) on each return value, which cascades into tb_S.realize()

So in this entire model register_buffer fires exactly once, on tb_S. That single call wraps Reduction_S in a ComputedBuffer and registers it as buf0. None of Pointwise_A, Pointwise_B, Pointwise_M ever get a bufN — they live entirely as bound methods inside Reduction_S.inner_fn’s closure chain.

Concretely: the cascade for ((a+b)*relu(c)).sum(-1)

When Stage 2 (or Stage 6) calls Reduction_S.inner_fn(red_idx), the captured-closure tree unwinds top-down:

Reduction_S.inner_fn(red_idx)            # the outermost call
  → Pointwise_M.inner_fn(red_idx)        # via the captured loader from make_loader()
      → ops.mul(
          Pointwise_A.inner_fn(red_idx),     # via captured loader for tb_M's input #0
          Pointwise_B.inner_fn(red_idx),     # via captured loader for tb_M's input #1
        )

Pointwise_A.inner_fn(red_idx):
  → ops.add(
      ops.load("arg0_1", red_idx),
      ops.load("arg1_1", red_idx),
    )

Pointwise_B.inner_fn(red_idx):
  → ops.relu(ops.load("arg2_1", red_idx))

Under the dependency-tracking handler (Stage 2), ops.load("arg0_1", ...) is recorded as a read of arg0_1, etc. The recorded read_writes for buf0’s SchedulerNode ends up:

reads  = {arg0_1, arg1_1, arg2_1}
writes = {buf0}

Under the Triton-codegen handler (Stage 6), the same calls produce:

tmp0 = tl.load(in_ptr0 + ..., rmask & xmask)    # arg0_1
tmp1 = tl.load(in_ptr1 + ..., rmask & xmask)    # arg1_1
tmp2 = tmp0 + tmp1                              # ops.add → tl + tl
tmp3 = tl.load(in_ptr2 + ..., rmask & xmask)    # arg2_1
tmp4 = triton_helpers.maximum(tmp3, 0)          # ops.relu
tmp5 = tmp2 * tmp4                              # ops.mul
... accumulator += tmp5 ...                     # the Reduction wrapping it

One realize call → one buf → one fused kernel covering all four ATen ops. The closure chain is the mechanism that makes this work without any explicit “fusion” pass having to recognize that add+relu+mul+sum can be merged: by the time anyone looks, it’s already one expression.

You can see this live with the deep_dive_ii_inductor_internals/run_walkthrough.py script — [hooks] register_buffer buf0 (ComputedBuffer) should be the only register_buffer line in the trace.

At this point the IR is a directed acyclic graph of Buffers. Each buffer has a name, a layout, and either a loop body (for ComputedBuffer) or a template binding (for TemplateBuffer).

Stage 2 — The Scheduler: where dependencies become read_writes

The next pass wraps each operation buffer in a BaseSchedulerNode to form a scheduler graph. From here on, the unit of work is a scheduler node; the unit of memory is a named buffer; and the unit of dependency information is a ReadWrites object attached to each scheduler node.

A BaseSchedulerNode carries a read_writes: ReadWrites attribute, which is a pair of ordered sets of Deps. There are several Dep flavors — MemoryDep (typed, indexed access), StarDep (whole-buffer read/write), WeakDep (ordering-only, doesn’t extend lifetime). The scheduler distinguishes them but for liveness purposes a name appearing in read_writes.reads or .writes keeps that buffer alive through this node.

How is read_writes populated? The fork is in SchedulerNode._compute_attrs:

if isinstance(self.node, ir.TemplateBuffer):
    self.set_read_writes(
        self.node.extract_read_writes(normalize=should_normalize)
    )
else:
    self.set_read_writes(
        dependencies.extract_read_writes(
            self._body, *self._sizes, normalize=should_normalize
        )
    )

This is the most important fork in the entire codegen pipeline.

Path A: ComputedBuffer — symbolic interpretation of the loop body

For a ComputedBuffer, the loop body is a Python callable (after simplify_and_reorder) that takes loop indices and emits ops via a virtualized handler. dependencies.extract_read_writes(...) runs this callable under a tracking handler that records every ops.load(name, idx) as a read and every ops.store(name, idx, ...) as a write.

Result: a ReadWrites that faithfully reflects every storage the kernel will touch, including any captured tensor pulled in via ops.load — because the only way a captured tensor enters the loop body is via a load.

Aliases get handled with one explicit walk: used_or_aliased_buffer_names() extends each read with get_inputs_that_alias_output() so a reinterpret_tensor(bufN, …) read keeps bufN alive. (is_fake WeakDeps are filtered out — they exist purely to constrain ordering.)

Path B: TemplateBuffer — the inputs-list shortcut

For a TemplateBuffer, the body is a Jinja template. There’s no Python loop body to interpret. Instead, the template buffer itself is asked: TemplateBuffer.extract_read_writes builds reads from _read_deps_from_inputs:

def _read_deps_from_inputs(self, normalize: bool) -> OrderedSet[dependencies.Dep]:
    """Build read dependencies from all inputs."""
    reads: OrderedSet[dependencies.Dep] = OrderedSet()
    for inp_raw in self.inputs:
        ...
        reads |= dependencies.extract_read_writes(
            dummy, inp.get_size(), (), normalize=normalize
        ).reads
    return reads

That for inp_raw in self.inputs: is the load-bearing line. Whatever the template kernel actually reads beyond its declared inputs list is invisible to the scheduler.

For a Triton template that’s mostly fine: epilogue / prologue fusion is implemented as separate scheduler nodes that retain their own loop bodies and read_writes, then get folded into the final fused node. The fold preserves their reads. Captures of “outer” tensors aren’t a thing for built-in Triton matmul-style templates.

set_read_writes, used_buffer_names, and friends

A few related plumbings keep read_writes accurate as the scheduler manipulates it:

• set_read_writes is the canonical setter; it always updates unmet_dependencies to read_writes.reads so dependency-resolution stays in sync.
• used_buffer_names flattens reads and writes to a name set.
• used_or_aliased_buffer_names extends with one alias hop.
• prune_deps() removes deps to ops that have been DCE’d or removed — it edits read_writes directly via set_read_writes(read_writes.remove_reads(to_remove)).

Everything downstream that asks “what does this node touch?” — fusion, last-use analysis, peak memory estimation, kernel signature emission — reads from these.

Stage 3 — Fusion: collapsing scheduler nodes

After basic scheduler graph construction, Inductor runs a fusion pass that collapses multiple SchedulerNodes into FusedSchedulerNodes. The pass is iterative: it tries every adjacent pair (vertical fusion = producer + consumer) or every pair sharing a parent (horizontal fusion = sibling consumers of the same producer), checks whether the backend can fuse them, and repeats until no more fusions are profitable.

Three predicates gate every fusion attempt:

Legality — can_fuse

The backend gates this. For Triton-eligible nodes, SIMDScheduling.can_fuse checks shape compatibility (matching iteration domains modulo broadcasting), reduction kinds (a split-scan can’t fuse with a reduction), tiling compatibility, and a long list of “would this be sound” predicates. The vertical case has its own can_fuse_vertical that additionally checks the producer-output-buffer can be inlined into the consumer.

Profitability — score_fusion_memory

Legality says yes/no; profitability ranks the survivors. score_fusion_memory estimates memory traffic saved — the bytes in buffers shared between the two nodes:

def score_fusion_memory(self, node1, node2):
    common_memory_deps = (
        (node1.read_writes.reads | node1.read_writes.writes) &
        (node2.read_writes.reads | node2.read_writes.writes)
    )
    return sum(self.dep_size_hint(dep) for dep in common_memory_deps)

Higher score = more bytes saved by fusion (because the shared buffer becomes register-resident instead of materialized in HBM). Pairs are sorted by this score so the most profitable fusions happen first. Note that this expression is a pure function of read_writes — every fusion priority decision reads off the same dep set computed in Stage 2.

Cycle safety — will_fusion_create_cycle

will_fusion_create_cycle runs a DFS to check if merging two nodes would introduce a circular dependency through other already-fused nodes. Fusion merges the ancestor sets of the two nodes, which can introduce cycles that didn’t exist before — so this check is essential before committing.

Post-fusion node types

The scheduler graph after fusion contains a small zoo of node types, all subclasses of BaseSchedulerNode:

Class	What it represents
SchedulerNode	Single un-fused op
FusedSchedulerNode	Multiple ops fused into one kernel
ExternKernelSchedulerNode	External library call (cuBLAS, cuDNN) — never fused
NopKernelSchedulerNode	No-op (alias, view) — emits no kernel
ForeachKernelSchedulerNode	Fused foreach ops (multi-tensor apply)

FusedSchedulerNode.set_last_usage does a two-level pass: first the global set_last_usage on the fused node (so other nodes see the fused boundary), then a reverse pass over the inner snodes to compute per-snode last-use. The inner pass is what enables Triton kernels to drop a register the moment a value is no longer needed inside the kernel body.

Stage 4 — Memory planning and last-use analysis

This is where the pipeline turns “node N reads buffer B” into “emit del B at line K of the wrapper output.”

compute_last_usage

The function is short. From scheduler.py:6811:

def compute_last_usage(self) -> None:
    """
    Populate node.last_usage recursively (also for the nodes within a
    FusedSchedulerNode)
    """
    future_used_buffers = OrderedSet(V.graph.get_output_names())

    for node in reversed(self.nodes):
        node.set_last_usage(future_used_buffers, self.mutation_real_name)
        future_used_buffers.update(node.last_usage)

The recipe in plain English: walk in reverse execution order; seed the “live set” with whatever the function returns (those have to stay live to the end); for each node, compute last_usage = used_or_aliased_buffer_names(self) - future_used_buffers; then add those names to future_used_buffers so earlier nodes see them as live.

set_last_usage implements the per-node piece:

def set_last_usage(self, future_used_buffers, mutation_real_name) -> None:
    used_buffers = self.used_or_aliased_buffer_names()
    used_buffers = OrderedSet(mutation_real_name.get(k, k) for k in used_buffers)
    self.last_usage = used_buffers - future_used_buffers

The mutation_real_name step renames in-place mutated buffers to their canonical name so a kernel that mutates buf3 in place isn’t treated as keeping buf2 (the pre-mutation alias) alive past its real death.

Peak memory and reuse planning

Beyond last-use, torch/_inductor/memory.py builds memory-planning information. assign_memory_planning_info_for_scheduler_buffers annotates every buffer with its size and, importantly, a list of “successors who use this buffer.” estimate_peak_memory_allocfree walks the schedule with this info, simulating allocations and frees, to produce the peak-memory estimate that drives reuse planning and topological re-ordering (topological_sort_lpmf).

Reuse is the optimization where, when buffer A dies on the same node where buffer B is born, the wrapper rebinds B’s name to A’s storage instead of allocating fresh memory: bufB = bufA; del bufA # reuse. This whole pass — what’s reusable, what reuse saves, where to put del lines — runs on top of last_usage, which runs on top of read_writes. Every layer in this stack inherits the dep set computed in Stage 2.

Stage 5 — Wrapper codegen: the visible artifact

Now we get to the Python file you can actually read in /tmp/torchinductor_*/.../output_code.py.

GraphLowering.codegen hands control to a PythonWrapperCodegen (or CppWrapperCpu for AOT-Inductor C++ output). The wrapper iterates the scheduler in execution order via Scheduler._codegen. For each node, it asks the relevant backend to codegen the kernel (Stage 6 below), then emits:

• Allocation lines. bufN = empty_strided_cuda(size, stride, dtype) if the buffer needs fresh storage, or bufN = reinterpret_tensor(bufM, …) for views, or bufN = bufM; del bufM # reuse for reuse rebinds.
• The kernel call. <kernel_name>.run(arg, arg, …, stream=streamN). The argument list comes directly from the kernel’s emitted signature (Stage 6).
• Free lines. One del per name in the node’s last_usage, mediated by codegen_free:

def codegen_free(self, buffer):
    name = buffer.get_name()

    if isinstance(buffer, (ir.InputBuffer, ir.TorchBindObject)):
        self.writeline(FreeLine(self, buffer))   # always del
        return

    if isinstance(buffer.get_output_spec(), ir.CommBufferLayout):
        self.writeline(FreeIfNotReusedLine(self, buffer, comm_buffer=True))
        return

    if not self.can_reuse(buffer):
        return
    self.freed.add(name)
    self.writeline(FreeIfNotReusedLine(self, buffer))

Two flavors:

• FreeLine for graph inputs and TorchBind objects — always emit del argN_1. The wrapper can’t repurpose input storage as scratch, so there’s no “maybe don’t free” path.
• FreeIfNotReusedLine for internal bufN — the wrapper waits until rendering time to decide. If a later allocation has been planned to reuse this buffer’s storage (reuse decision happens during the scheduler walk), no free line is emitted; otherwise, del bufN.

There’s also a sweep at kernel boundaries: Scheduler.free_buffers drains a buffer_names_to_free queue, ensuring nothing lingers past its computed death.

A representative slice of generated wrapper output:

def call(args):
    arg0_1, arg1_1, arg2_1 = args
    args.clear()
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda(...)
        triton_poi_fused__to_copy_view_0.run(arg0_1, buf0, ...)
        buf1 = empty_strided_cuda(...)
        triton_red_fused_sum_1.run(buf0, arg1_1, buf1, ...)
        del buf0
        del arg0_1
        return (buf1,)

Every del line comes from last_usage for the node above it; every bufN = ... comes from an allocation request emitted during scheduler walk; every <kernel>.run(...) argument list comes from the kernel’s signature (next stage). All of which derive from the dep set computed in Stage 2.

Stage 6 — Triton kernel codegen: where the dep set becomes tl.load

This is where the abstraction becomes physical. For Triton-eligible scheduler nodes, SIMDScheduling drives kernel emission via its codegen_node entry point, which assembles a TritonKernel (a subclass of SIMDKernel) and calls codegen_kernel.

The virtualized ops pattern

Before diving into mechanics, the key abstraction. Each IR node’s computation is expressed as an inner_fn that calls virtualized ops — ops.load, ops.store, ops.add, ops.relu, etc. These ops are not real functions; they dispatch through V.set_kernel_handler to whatever handler is currently installed:

# The same inner_fn produces different output depending on the installed handler:
def inner_fn(index):
    tmp0 = ops.load("buf0", index)
    tmp1 = ops.relu(tmp0)
    return tmp1

# During Triton codegen:  ops.load → "tl.load(in_ptr0 + xindex, xmask)"
# During C++ codegen:     ops.load → "Vectorized<float>::loadu(in_ptr0 + x0, 16)"
# During analysis:        ops.load → increments a read byte counter

This pattern is what lets the same IR be lowered to Triton, C++ vectorized loops, or pure analysis passes (memory traffic estimation, dep extraction in Stage 2) without rewriting the IR. The current handler is a context-managed stack — you push a TritonKernel’s overrides for the duration of its body emission, then pop.

The kernel signature is a function of read_writes

The Triton kernel’s parameter list is built by walking the fused node’s reads (input pointers), writes (output pointers), and any auxiliary scalars or sizes. The wrapper then matches that signature when emitting <kernel>.run(...). The same read_writes from Stage 2 that drove last-use analysis also drives signature construction — every tl.load(in_ptr_N + ...) in the generated kernel corresponds to one MemoryDep in read_writes.reads.

Tiling and indexing

SIMDKernel.codegen_indexing takes a sympy index expression and rewrites it into an “indexing variable” that the kernel can use; this is where strided loads turn into something like tl.load(in_ptr0 + (x0 + 64 * x1)). SIMDScheduling.select_tiling chooses how to split the iteration space across XBLOCK / RBLOCK (and sometimes YBLOCK) — it decides the kernel’s pid axes.

Tiling decisions feed back into masks: a kernel processing numel = 1234 with XBLOCK = 1024 needs xmask = xindex < xnumel to avoid out-of-bounds loads on the tail block. The mask emission piggybacks on indexing.

codegen_body and codegen_kernel

TritonKernel.codegen_body walks the loop body once more, this time emitting Triton ops (tl.load, tl.store, tl.sum, tl.where, etc.) into an IndentedBuffer. Reductions wrap their accumulators in RBLOCK loops; pointwise ops are flat. CSE (common subexpression elimination) deduplicates intermediate computations across snodes within the same kernel via the CSE class — tmp0 = tl.load(...) only happens once per (name, index) pair.

TritonKernel.codegen_kernel wraps it all into a @triton.jit-decorated function: typed argdefs come from the KernelArgs bookkeeper, jit_lines emits the @triton_heuristics.pointwise(...) decorator (with grid metadata, num_warps, num_stages, etc.), and the body is the buffer from codegen_body. The result is a string spliced into the generated Python file with a name like triton_poi_fused__to_copy_add_clone_view_NN.

A peek at what comes out for a simple (a + b) * c.relu():

@triton_heuristics.pointwise(
    size_hints={'x': 32768},
    filename=__file__,
    triton_meta={...},
    inductor_meta={...},
)
@triton.jit
def triton_poi_fused__to_copy_add_clone_view_3(
    in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr,
):
    xnumel = 26214400
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp1 = tl.load(in_ptr1 + (x0 % 4096), xmask)
    tmp2 = tmp0 + tmp1
    tl.store(out_ptr0 + (x0), tmp2, xmask)

Everything in here — in_ptr0, in_ptr1, out_ptr0 — corresponds to a name in the wrapper’s argument list, which corresponds to a buffer in read_writes, which corresponds to a name allocated by register_buffer back in Stage 1.

For TemplateBuffer nodes the path is different: SIMDScheduling.codegen_template renders a Jinja template and stitches it into the generated file.

Aside: how flex_attention threads the pipeline as a TemplateBuffer

The whole post has been about Path A — the Pointwise/Reduction closure path that dominates for ATen ops. But the side note in Stage 1 mentioned a parallel path: matmul, conv, and flex_attention go through TritonTemplateBuffer instead. That path is shaped differently at every stage. Worth walking through end-to-end with flex_attention as the example, because it’s where most of the practical interest in templates lives today.

Stage 1 — registered at construction, always realized

flex_attention isn’t a plain ATen op; it’s a higher-order op (torch.ops.higher_order.flex_attention) whose argument list includes two FX subgraphs — subgraph (the user’s score_mod) and block_mask (which contains the user’s mask_mod). Its lowering lives at torch/_inductor/kernel/flex/flex_attention.py:

@register_lowering(torch.ops.higher_order.flex_attention, type_promotion_kind=None)
def flex_attention(query, key, value, subgraph, block_mask, scale,
                   kernel_options, score_mod_other_buffers, mask_mod_other_buffers):
    ...

The template itself is a module-level constant — flex_attention_template:

flex_attention_template = TritonTemplate(
    name="flex_attention",
    grid=flex_attention_grid,
    source=load_flex_template("flex_attention")
         + load_flex_template("utilities")
         + load_flex_template("common"),
    always_freeze_layout=True,
)

TritonTemplate is a KernelTemplate subclass that wraps a Jinja string. The body is not a Python inner_fn; it’s text with Jinja substitutions like ,, ``. The lowering then calls flex_attention_template.maybe_append_choice(...):

flex_attention_template.maybe_append_choice(
    choices=choices,
    input_nodes=[query, key, value, logsumexp, max_scores,
                 kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices],
    layout=layout,
    subgraphs=[subgraph_buffer, mask_graph_buffer],
    mutated_inputs=[logsumexp, max_scores],
    call_sizes=query.get_size(),
    **cur_kernel_options,
)

This constructs a TritonTemplateBuffer, and TemplateBuffer.__init__ calls V.graph.register_buffer(self) immediately. No realize() step is needed and the closure path doesn’t apply. From the moment a flex_attention lowering returns, there’s a bufN for it on V.graph.buffers.

The Pointwise closure mechanism is deliberately bypassed because the kernel body is a Jinja template, not a Python callable. There is nothing to inline a make_loader into — Inductor can’t symbolically interpret a 200-line flash_attention.j2.

Stage 2 — Path B in _compute_attrs

When the scheduler wraps the TritonTemplateBuffer in a SchedulerNode, _compute_attrs takes the isinstance(self.node, ir.TemplateBuffer) branch — Path B. Reads come from TemplateBuffer.extract_read_writes, which delegates to _read_deps_from_inputs iterating self.inputs. For our flex_attention call, that’s the 9 named inputs (q, k, v, lse, max_scores, four block_mask buffers).

This is exactly the underreport-prone shortcut from the closure section — read deps come from a declared list, not from symbolic interpretation. For a Triton flex template this works out OK because score_mod / mask_mod captures are passed through score_mod_other_buffers / mask_mod_other_buffers which the flex_attention lowering threads into inputs_for_autotuning separately (the autotuner sees them, even if the template buffer’s named-input list does not).

Stage 6 — SIMDScheduling.codegen_template and the Jinja render

At codegen time, SIMDScheduling.codegen_template is the entry for any TemplateBuffer node. It instantiates a TritonTemplateKernel and calls render:

def render(self, template, kwargs, ...):
    template_env = {
        fn.__name__: (...)
        for fn in [self.def_kernel, self.size, self.stride,
                   self.store_output, self.load_input, self.make_load,
                   self.modification, self.gen_argdefs, self.gen_defines, ...]
    }
    return PartialRender(template.render(**template_env, **kwargs), ...)

Each method on TritonTemplateKernel (def_kernel, store_output, load_input, modification, …) becomes a callable inside the Jinja namespace. When the template hits , that fires `self.def_kernel(...)`, which emits the `@triton.jit` signature. emits a tl.store(...). The template is essentially programming the kernel codegen by composing these emitters in a fixed structure (load Q tile → load K tile → matmul → optional score_mod → softmax → matmul with V → store).

The interesting one is **** — this is where the user's `score_mod` and `mask_mod` get inlined. Each modification call corresponds to a subgraph buffer from `subgraphs=[...]` and emits the subgraph's body (which *is* a small `Pointwise`-style closure!) into the kernel at exactly that point. So inside the template you have ordinary Triton code like `score = ...`, then a placeholder that expands to the user’s score_mod body inlined at that line — turning one cross-template hook into one nested chunk of generated Triton.

Prologue / epilogue fusion

The last piece. The scheduler can absorb adjacent pointwise nodes into a TemplateBuffer kernel via fusion. _codegen_single_template and codegen_template_body walk the prologue and epilogue scheduler nodes (preceding and following pointwise consumers/producers) and weave their inner_fns into the same kernel via load_input (prologue) and store_output (epilogue) hooks. So a sequence like flex_attention(q, k, v) * scale + bias can fuse the trailing *scale + bias into the same kernel without writing the attention output to HBM.

This is where Path A and Path B meet: the template kernel is fixed-shape, but its prologue / epilogue slots accept arbitrary Pointwise chains (i.e. closure-form subgraphs) that get inlined via virtualized ops. The read_writes for the resulting fused kernel union both: the named inputs from the template, plus everything the prologue / epilogue inner_fn closures touch.

Summary: how the two paths differ

Aspect	Path A (Pointwise / Reduction)	Path B (TemplateBuffer, e.g. flex_attention)
Kernel body	Python inner_fn closure	Jinja template string
Realization	Lazy via TensorBox.realize() (only when forced)	Eager — register_buffer in TemplateBuffer.__init__
read_writes source	Symbolic interpretation of inner_fn (Stage 2 Path A)	Declared self.inputs list (Stage 2 Path B)
Codegen entry	SIMDScheduling.codegen_node	SIMDScheduling.codegen_template
Op-by-op IR	Yes (CSE deduplicates expressions)	No — kernel is a fixed structure with hook calls
Fusion	Driven by adjacent-snode merge in FusedSchedulerNode	Prologue / epilogue absorption of neighbors
Subgraph inlining	N/A	`` hooks render score_mod / mask_mod inline

Both paths produce a bufN line in the wrapper, a triton_*.run(...) call, and (eventually) a tl.load/tl.store body. They get there through completely different machinery, and you can tell which path a given kernel took from the kernel name alone — triton_red_fused_* is Path A, triton_tem_fused_* (or the cutedsl-named variants) is Path B.

Closing

The Inductor pipeline is a chain of layers, each consuming the previous layer’s view of “what does this kernel touch?”:

• Stage 1 (lowering) registers buffers and assigns names.
• Stage 2 (scheduler) computes a ReadWrites per node — symbolically for ComputedBuffer, by declared inputs for TemplateBuffer. The fork.
• Stage 3 (fusion) unions / rewrites those sets across fused groups, gated by can_fuse (legality), score_fusion_memory (profitability), and will_fusion_create_cycle (safety).
• Stage 4 (memory planning) turns the union into per-node last_usage and reuse decisions.
• Stage 5 (wrapper) renders allocations, kernel calls, and dels.
• Stage 6 (kernel codegen) uses the same dep set to build kernel signatures via the virtualized-ops dispatch pattern.

Buffer lifetimes — when each bufN is born, used, reused, and del‘d — are the most directly observable consequence of all of this. Read any inductor-generated output_code.py and you can reason backward from a del bufN line through last_usage → read_writes → the underlying IR node, and end up at the exact ATen op in the source FX graph that produced it.

If you want to go deeper, the parts I cut for length: how MultiOutputLayout templates compose with the _read_deps_from_inputs shortcut; the role of mutation_real_name and how in-place ops work without breaking liveness; how MemoryPlanning’s peak-memory estimates feed back into topological_sort_lpmf to reorder the schedule; how the C++ wrapper path (AOTInductor) differs; and the foreach / split-scan / persistent-reduction kernel templates that have their own scheduling quirks.

The code is messier than the diagram, but the diagram is the spine.