How compilation works

Open in Colab

The previous tutorial showed how to compile a model and call it — neml2-compile, the .pt2 artifact, the drop-in stub. This one opens the box: what actually happens when a model is compiled, shown on a small example where the compiler’s optimization is visible in the output.

neml2-compile works in two stages, both powered by PyTorch:

  1. Capture — torch.export traces the eager model into a backend-agnostic graph of tensor operations (an FX graph).

  2. Lower — AOTInductor takes that graph, optimizes it, and generates + compiles a native kernel (C++ on CPU, Triton on GPU), packaged into the .pt2.

We’ll watch a model go through both stages and read the graph and the generated code at each end.

A model with a shared subexpression

To make the optimization visible we need an example where the same quantity is computed more than once. Three first-order rates sharing a denominator do the job:

\[ \dot{a} = \frac{a - a_n}{t - t_n}, \qquad \dot{b} = \frac{b - b_n}{t - t_n}, \qquad \dot{c} = \frac{c - c_n}{t - t_n}, \]

where the subscript \(n\) is the previous time step. Each equation is a ScalarVariableRate; composing the three gives a model that takes \(a, a_n, b, b_n, c, c_n, t, t_n\) and returns \(\dot a, \dot b, \dot c\). The denominator \(t - t_n\) is shared across all three — a redundancy the compiler should be able to remove.

%%writefile input.i
[Models]
  [eq1]
    type = ScalarVariableRate
    variable = 'a'
  []
  [eq2]
    type = ScalarVariableRate
    variable = 'b'
  []
  [eq3]
    type = ScalarVariableRate
    variable = 'c'
  []
  [eq]
    type = ComposedModel
    models = 'eq1 eq2 eq3'
  []
[]

Writing input.i

Load it the usual way — at this point it’s an ordinary eager model, a torch.nn.Module composed from the three leaves:

import torch

torch.set_default_dtype(torch.float64)

import neml2

model = neml2.load_model("input.i", "eq")
print("inputs :", list(model.input_spec))
print("outputs:", list(model.output_spec))

inputs : ['a', 'a~1', 't', 't~1', 'b', 'b~1', 'c', 'c~1']
outputs: ['a_rate', 'b_rate', 'c_rate']

Stage 1 — capture: torch.export → the FX graph

The first thing neml2-compile does is trace the model with torch.export, recording every tensor operation into an FX graph. Tracing needs example inputs with a representative dtype and batch shape — one typed tensor per declared input (a batch of 2 here, so the batch axis is a real dimension):

from neml2.types import Scalar
from torch.export import export

# One example input per declared input, in `input_spec` order.
example_inputs = tuple(Scalar(torch.full((2,), float(i + 1))) for i in range(len(model.input_spec)))

exported = export(model, example_inputs, strict=False)

# Print just the captured forward body (the tensor operations).
graph_src = exported.graph_module.print_readable(print_output=False)
for line in graph_src.splitlines():
    if "aten." in line or line.strip().startswith(("def forward", "return")):
        print(line.rstrip())

    def forward(self, inputs_0_data: "f64[2]", inputs_1_data: "f64[2]", inputs_2_data: "f64[2]", inputs_3_data: "f64[2]", inputs_4_data: "f64[2]", inputs_5_data: "f64[2]", inputs_6_data: "f64[2]", inputs_7_data: "f64[2]"):
        sub: "f64[2]" = torch.ops.aten.sub.Tensor(inputs_2_data, inputs_3_data)
        sub_1: "f64[2]" = torch.ops.aten.sub.Tensor(inputs_0_data, inputs_1_data);  inputs_0_data = inputs_1_data = None
        div: "f64[2]" = torch.ops.aten.div.Tensor(sub_1, sub);  sub_1 = sub = None
        sub_2: "f64[2]" = torch.ops.aten.sub.Tensor(inputs_2_data, inputs_3_data)
        sub_3: "f64[2]" = torch.ops.aten.sub.Tensor(inputs_4_data, inputs_5_data);  inputs_4_data = inputs_5_data = None
        div_1: "f64[2]" = torch.ops.aten.div.Tensor(sub_3, sub_2);  sub_3 = sub_2 = None
        sub_4: "f64[2]" = torch.ops.aten.sub.Tensor(inputs_2_data, inputs_3_data);  inputs_2_data = inputs_3_data = None
        sub_5: "f64[2]" = torch.ops.aten.sub.Tensor(inputs_6_data, inputs_7_data);  inputs_6_data = inputs_7_data = None
        div_2: "f64[2]" = torch.ops.aten.div.Tensor(sub_5, sub_4);  sub_5 = sub_4 = None
        return (div, div_1, div_2)

Read the graph against the input order printed above: inputs_2_data and inputs_3_data are \(t\) and \(t_n\), so every aten.sub.Tensor(inputs_2_data, inputs_3_data) is the denominator \(t - t_n\). It appears three times (sub, sub_2, sub_4) — once per leaf. The capture is faithful but unoptimized: each ScalarVariableRate computed its own copy of the shared denominator, exactly as the eager Python ran. This is the device-agnostic starting point; the optimization happens next.

Stage 2 — lower: AOTInductor → native code

neml2-compile hands the FX graph to AOTInductor, which runs a series of optimization passes (common-subexpression elimination, fusion, layout choices, dead-code elimination, vectorization, …), generates kernel source for the target device, compiles it, and packages the result. Let’s run the real compile (the same command from the previous tutorial) and then read what it generated:

!neml2-compile input.i --model eq

Compiled 'eq' (composed, fully baked, no derivative graphs (use -d to compile jvp/jacobian)) for [cpu] -> aoti/eq
  stub: aoti/eq_aoti.i
  cpu: aoti/eq/cpu/eq_meta.json

A .pt2 is just a zip archive. Inside it are the generated C++ source (a wrapper.cpp that handles allocation and the calling convention, and a kernel.cpp holding the actual math) and the compiled .so the runtime loads:

import zipfile

pt2 = "aoti/eq/cpu/eq.pt2"
with zipfile.ZipFile(pt2) as z:
    for name in z.namelist():
        if name.endswith((".cpp", ".so")):
            print(name.split("/")[-1])

cr2pvutsngf23znknbuqrxhc7qxzmxpsxxlplo2axi5r36cr2rvf.wrapper.cpp
c4x6rpqev7wuvvctheyorfxcv7fh7oifc3aycgnipojdpk3xk7l2.kernel.cpp
cr2pvutsngf23znknbuqrxhc7qxzmxpsxxlplo2axi5r36cr2rvf.wrapper.so

Now the payoff — read the compute kernel and look at the arithmetic the compiler actually emitted:

with zipfile.ZipFile(pt2) as z:
    kernel = next(n for n in z.namelist() if n.endswith("kernel.cpp"))
    src = z.read(kernel).decode()

# The fused kernel's name, then its arithmetic. The kernel has two loops (a
# vectorized main loop and a masked tail) with identical math, so we print the
# first one and stop once the arithmetic starts repeating.
print(next(line.strip() for line in src.splitlines() if "cpp_fused" in line and "void" in line))
print("...")
seen = set()
for line in src.splitlines():
    s = line.strip()
    if s.startswith("auto tmp") and (" - tmp" in s or " / tmp" in s):
        if s in seen:
            break
        seen.add(s)
        print("   ", s)

extern "C"  void  cpp_fused_div_sub_0(const double* in_ptr0,
...
    auto tmp2 = tmp0 - tmp1;
    auto tmp5 = tmp3 - tmp4;
    auto tmp6 = tmp2 / tmp5;
    auto tmp9 = tmp7 - tmp8;
    auto tmp10 = tmp9 / tmp5;
    auto tmp13 = tmp11 - tmp12;
    auto tmp14 = tmp13 / tmp5;

Three things happened here, all visible in those seven lines:

  • Common-subexpression elimination. tmp5 = tmp3 - tmp4 is the denominator \(t - t_n\), computed once and reused by all three divisions (tmp6, tmp10, tmp14). The threefold redundancy from the FX graph is gone.

  • Fusion. All three rates live in a single function (cpp_fused_div_sub_0), one loop over the batch — not three separate kernels writing intermediates to memory and reading them back.

  • Vectorization. The arithmetic is on at::vec::Vectorized SIMD registers, so several batch entries are processed per instruction.

This is the same common-subexpression optimization NEML2’s older just-in-time engine performed on its trace — but here it lands as real, vectorized, compiled machine code rather than an optimized interpreter graph.

The IR in between

Between the FX graph and the C++, Inductor builds its own intermediate representation and makes the fusion and scheduling decisions there. It’s verbose and internal, so we don’t reproduce it here — but if you’re curious you can dump the entire pipeline (the FX graph, the pre- and post-fusion IR, and the final generated code) by setting an environment variable before compiling:

TORCH_COMPILE_DEBUG=1 neml2-compile input.i --model eq

This writes a torch_compile_debug/ folder; ir_post_fusion.txt shows the three operations collapsed into a single FusedSchedulerNode, and output_code.cpp is the kernel we read above.

What changes when we compile for the GPU

This notebook compiled for the CPU, so the kernel above is C++. Targeting CUDA (neml2-compile --device cuda, which additionally needs nvcc) reuses the same captured FX graph but lowers it differently. A few things worth knowing:

  • Same capture, different codegen. Stage 1 (the FX graph) is identical — capture is device-agnostic. Only Stage 2 changes: Inductor emits Triton GPU kernels instead of C++, then the same optimization passes (CSE, fusion, vectorization) run against the GPU cost model.

  • Fusion matters more — kernel launch cost. Every GPU kernel launch carries a fixed overhead (microseconds). Fusing the three rates into one kernel means one launch instead of several; at small-to-moderate batch sizes, where there isn’t enough arithmetic to hide it, launch overhead can dominate — so fewer, larger kernels pay off more than on CPU.

  • Fewer temporaries, less memory traffic. GPUs are usually memory-bandwidth-bound. Fusion keeps intermediates like \(t - t_n\) in registers instead of writing them out to global memory and reading them back; that means fewer temporary buffers and less round-tripping to HBM, which is frequently the real bottleneck rather than the arithmetic itself.

  • Autotuning. At compile time Inductor autotunes each Triton kernel — block sizes, tiling, the number of warps — by benchmarking candidate configurations for the target shape and keeping the fastest. It’s a one-time compile cost baked into the artifact. CUDA-graph capture (replaying a whole launch sequence as one unit) is also available for launch-bound graphs.

  • Same artifact contract. The result is still a .pt2, just holding a CUDA .so; --device cpu cuda builds both and the loader picks the folder matching the runtime device (see Compiled models). Device and dtype are pinned at compile time.

Recap

When we run neml2-compile, the model travels:

eager nn.Module → torch.export (FX graph) → AOTInductor (optimize → generate → compile native kernels) → .pt2 package

The first stage is a faithful, unoptimized capture; the second is where the redundancy is removed and the math becomes fast device code.

Where to go next

  • Compilation pipeline — the full compilation pipeline reference, including how implicit (Newton-solve) models are split into segments.

  • AOTI packages — the on-disk layout and metadata contract of a .pt2 artifact.

  • Compiled models — back to the user-facing view: loading and running a compiled model.