A machine learning compiler translates tensor programs into executable code for the target hardware—here, GPU kernels—while optimizing execution and preserving the required numerical behavior.

One line of PyTorch can become very different compiler graphs. For a linear layer, a machine learning compiler can keep matrix multiplication explicit, expand it into multiplication and reduction, or retain several equivalent forms. Each choice changes what later stages can optimize directly and what structure they must rediscover before selecting GPU code.

To see how compilers make these decisions under the hood, we trace a toy example of one linear operation through PyTorch Inductor, MAX, TinyGrad, and Luminal, all targeting an NVIDIA H100 GPU. Part 1 examines each compiler's intermediate representation (IR)—its operations, types, and value dependencies—before kernel boundaries are frozen. Part 2 (forthcoming) continues through storage allocation, implementation selection, kernel scheduling, code generation, and execution.

The PyTorch line below specifies the computation. Our comparison setup builds it through each system's frontend, using the same input shapes and numerical contract. The excerpts are saved captures from selected semantic stages.

Python
# FP16 PyTorch tensors; contiguous row-major inputs:
# x: [1024, 768], weight: [1024, 768] ([out_features, in_features]),
# b: [1024] -> y: [1024, 1024]
y = torch.matmul(x, weight.T) + b
A directed acyclic graph: weight feeds transpose; x and the transposed weight feed matrix multiplication; the projection proj and bias b feed addition to produce y.
Figure 1.1. The source computation as a directed acyclic graph (DAG). Arrows show value dependencies. The operation nodes do not specify GPU kernel boundaries, and the transpose does not necessarily copy data. The FP16 result proj is rounded before bias addition.

One Program and One Numerical Contract

PyTorch's nn.Linear(768, 1024) defines the transformation y = x @ weight.T + b, with weight shaped [out_features, in_features] = [1024, 768]. Our toy workload uses that weight convention with contiguous row-major inputs and explicitly constructs transpose, MatMul, and add operations. It is not a trace of nn.Linear.

Each of the 1,024 input rows has 768 features and produces 1,024 output features. The [1024, 1024] projection proj sums products over the shared input dimension K = 768; this sum is the contraction. In the formula, m indexes input rows, n indexes output features, and k indexes input features. We name the intermediate proj even though the source expression stays on one line. FP16 and FP32 mean 16-bit and 32-bit floating point.

projm,n=FP16(∑k=0767xm,k·weightn,k)

Adding the [1024] bias b is an epilogue: every row uses the same b[n] for output column n.

In eager FP16 execution, the separate torch.matmul returns the FP16 tensor proj before + b runs. We preserve that rounding boundary and additionally require FP32 matrix arithmetic:

  1. FP32 matrix arithmetic. Products accumulate in FP32 without first rounding each product to FP16. The contract allows the order of summation to vary.
  2. FP16 projection rounding. The matrix result proj is rounded to FP16 before bias addition, and the final output y is FP16.

Rounding before bias can change the answer, so it constrains which rewrites are legal. These are requirements of this comparison; the nn.Linear API does not promise an FP16 intermediate before bias addition. A fused kernel must preserve our rounding step, but it can keep proj in registers instead of storing it in memory.

With that, we can compare what each representation keeps explicit, what that enables, and what later stages must still decide.

Inductor Retains MatMul and Explicit Casts

We wanted TorchInductor to share similar abstractions to PyTorch eager to allow support of all of PyTorch, with a thin translation layer.

— PyTorch 2 paper, §4.1

Inductor is the default compilation backend for torch.compile(). Its design principles prioritize broad PyTorch coverage: tensor strides, views, and in-place updates must fit the compiler's model. A view changes how tensor elements are indexed while sharing the underlying storage. Mirroring PyTorch's tensor and storage abstractions reduces the translation needed to support those behaviors. Implementing the compiler in Python also makes it easier for PyTorch users to extend. In our example, the transposed weight is a concrete reason to preserve view semantics while optimizing the surrounding computation.

The captured FX graph represents tensor operations and their dependencies. It mixes ATen tensor operators (torch.ops.aten.*) with smaller PrimTorch primitives (torch.ops.prims.*), such as dtype conversions. The compiled FX graph retains aten.mm while surrounding the bias epilogue with explicit primitive casts:

Explore the Inductor graph

Bold labels describe the computation; monospace names match the IR. Select a node to inspect its inputs and captured code.

Operation Tensor / input view
MatMul
IR name(s)
aten.mm → mm
Inputs
Input x (arg1_1), Transpose (aten.permute → permute)
Shape / type
[1024, 1024] · FP16

Retains matrix multiplication as one ATen operation. Its output is the FP16 projection; accumulation precision is controlled by the compiler settings.

Captured lines 7
        mm: "f16[1024, 1024]" = torch.ops.aten.mm.default(arg1_1, permute);  arg1_1 = permute = None

The matrix result is already FP16. The later FP32 cast widens that rounded value before bias addition. Arrows show value dependencies, not kernel launches.

Full captured excerpt
PyTorch FX
permute: "f16[768, 1024]" = torch.ops.aten.permute.default(arg0_1, [1, 0]);  arg0_1 = None
mm: "f16[1024, 1024]" = torch.ops.aten.mm.default(arg1_1, permute);  arg1_1 = permute = None
convert_element_type_2: "f32[1024, 1024]" = torch.ops.prims.convert_element_type.default(mm,
    torch.float32);  mm = None
convert_element_type_3: "f32[1024]" = torch.ops.prims.convert_element_type.default(arg2_1,
    torch.float32);  arg2_1 = None
add: "f32[1024, 1024]" = torch.ops.aten.add.Tensor(convert_element_type_2,
    convert_element_type_3);  convert_element_type_2 = convert_element_type_3 = None
convert_element_type_4: "f16[1024, 1024]" =
    torch.ops.prims.convert_element_type.default(add, torch.float16);  add = None

Here arg0_1, arg1_1, and arg2_1 are weight, x, and b; each node carries static shape and dtype metadata ("f16[1024, 1024]"), and = None clears references after their last use in this graph listing. aten.mm returns the FP16 projection. The two FP32 casts inserted by our comparison setup widen that already-rounded result and the bias for addition; the final cast returns FP16. This preserves the projection's rounding boundary while performing the addition in FP32.

Keeping aten.mm gives later stages an explicit contraction to match to a library call or matrix template, while the visible bias arithmetic can participate in pointwise optimization. Fusion must connect those two representations while preserving the projection rounding.

MAX Preserves Explicit Matrix Semantics

MAX offers a layered architecture that lets you work at the level of abstraction that best fits your needs.

— MAX Python API overview

Like Inductor, MAX keeps MatMul explicit. The difference in this capture is how much detail surrounds it: MAX retains a compact transpose–MatMul–add graph, while the Inductor capture exposes the epilogue's widening and narrowing casts.

Modular's kernel-generation design rationale emphasizes established compiler techniques, predictable lowering, and extensibility without changing compiler source. The goal is to let kernel authors support new algorithms and hardware without waiting for compiler developers to encode every optimization. Parameterized kernel implementations carry that expertise; author-supplied constraints can limit the implementation search.

MAX exposes a graph-building API and runtime for model authors, while Mojo gives kernel authors a general-purpose language with explicit low-level control beyond the restrictions of an embedded domain-specific language (DSL). MAX lowers graphs through successive compiler representations. Its input graph uses MLIR (Multi-Level Intermediate Representation). Modular's rmo dialect supplies the operation and type vocabulary shown below. Names such as %2 identify values, each defined once—SSA, or Static Single Assignment:

MLIR
%2 = rmo.mo.transpose(%arg1, %1) : (!mo.tensor<[1024, 768], f16, gpu:0>, !mo.tensor<[2],
    si64>) -> !mo.tensor<[768, 1024], f16, gpu:0>
%3 = rmo.matmul(%arg0, %2) : (!mo.tensor<[1024, 768], f16, gpu:0>, !mo.tensor<[768, 1024],
    f16, gpu:0>) -> !mo.tensor<[1024, 1024], f16, gpu:0>
%4 = rmo.add(%3, %arg2) : (!mo.tensor<[1024, 1024], f16, gpu:0>, !mo.tensor<[1024], f16,
    gpu:0>) -> !mo.tensor<[1024, 1024], f16, gpu:0>
MAX value flow: weight feeds transpose; x and the transposed weight feed MatMul; the FP16 projection and bias feed add. SSA names match the excerpt.
Figure 1.2. Read the MAX excerpt as a value graph: %2 is the transposed weight, %3 the projection, and %4 the output. Blue boxes are operations; white boxes are tensors. The permutation constant and full type signatures are omitted.

%arg0, %arg1, and %arg2 are x, weight, and b, and %1 is the [1, 0] permutation constant that swaps the two axes. The !mo.tensor<[shape], dtype, device> signatures record each input and output's shape, storage dtype (f16), and device (gpu:0). They do not specify the internal accumulation precision of rmo.matmul.

This is an example of progressive lowering: keep useful structure until a stage can map it to the target. Here MatMul remains available for implementation selection; the graph alone does not reveal the tuning policy or kernel schedule.

TinyGrad Decomposes MatMul into UOps

Everything in Tensor is syntactic sugar around constructing a graph of UOps.

— TinyGrad developer documentation

TinyGrad takes a different starting point from both Inductor and MAX: the contraction is already decomposed in this semantic capture. Its design emphasizes a small, modifiable stack and a limited set of low-level operations for accelerator support. Expressing tensor APIs through shared primitives reduces the operation-specific machinery a backend must implement. Lazy graph construction lets the compiler see their composition before deciding what to execute.

The frontend builds a graph of UOps. A UOp is an IR node: op names the operation, src holds its inputs, and arg stores metadata. With no primitive matmul UOp, the frontend decomposes the contraction into 3D broadcast views ([1024, 1024, 768]), elementwise multiplication (Ops.MUL), and sum reduction (Ops.REDUCE) (semantic UOps):

TinyGrad UOps
UOp(Ops.ADD, arg=None, src=(
  UOp(Ops.CAST, arg=dtypes.f16, src=(
    UOp(Ops.REDUCE, arg=(Ops.ADD, 1), src=(
      UOp(Ops.PERMUTE, arg=(2, 0, 1), src=(
        UOp(Ops.CAST, arg=dtypes.f32, src=(
          UOp(Ops.MUL, arg=None, src=(
TinyGrad value flow: input and weight views feed multiply, FP32 cast, permutation and reduction, FP16 projection rounding, then bias addition.
Figure 1.3. Read the TinyGrad excerpt from inputs to output, reversing its root-first nesting. White boxes are inputs; blue boxes are operations. Input views and permutation/reduction are grouped. The virtual 3D product describes indexing, and the arrows do not specify GPU kernels.

The dump prints root-first from output y (Ops.ADD) inward through src=(...). Read Figure 1.3 in the opposite direction, from inputs to output. Movement ops reshape x (ParamArg(1)) to [1024, 1, 768] and permute/reshape weight (ParamArg(15)) to [1, 1024, 768]. Broadcasting aligns them over [m, n, k] = [1024, 1024, 768]: each virtual element is the product x[m,k] × weight[n,k]. These views change indexing without allocating a 3D buffer.

The graph casts the product to FP32, permutes k to the front ([768, 1024, 1024]), and sums that leading axis. In arg=(Ops.ADD, 1), 1 means one leading axis is reduced. The result is cast to FP16 before adding bias (ParamArg(17)).

In the pinned version, matmul(..., dtype=float32) delegates to dot, expressed as (x*w).sum(..., dtype=dtype). This is why CAST(dtypes.f32) wraps MUL in the dump (matrix frontend, reduction cast). The nesting records the logical expression; it does not by itself establish the precision of the executed multiplication.

Here the same indexing and fusion rewrites can work across views, contraction, and bias. Using a specialized matrix instruction then requires recognizing the contraction pattern again. The small vocabulary also leaves code generation, allocation, and execution support as target-specific work.

Luminal Retains Equivalent Expressions

Luminal tries to search every possible decision to give the compiler the flexibility to discover complex optimizations.

— Luminal README, “Search”

Like TinyGrad, Luminal starts from primitive operations. Its distinctive design goal is to use search to reduce reliance on hand-written performance heuristics. Equivalence rules describe which forms are legal; plan selection determines which available implementation is worth using. In this capture, a primitive contraction and recognized matrix forms remain available together.

The initial graph expresses matrix multiplication through 3D broadcast views, multiplication (LogicalMul), and reduction (LogicalReduceSum). Because its frontend computes in the operands' dtype, our comparison setup inserts explicit FP32 input casts before multiplication and an FP16 projection cast after reduction (logical program):

Egglog
(let v4 (LogicalCast v0 (F32)))
(let v5 (LogicalCast v3 (F32)))
(let v8 (LogicalMul v6 v7))
(let v9 (LogicalReduceSum v8 0))
(let v10 (LogicalCast v9 (F16)))
(let v13 (LogicalAdd v10 v12))

Each let binds a value. The inputs v0, v1, and v2 are x, weight, and b. The omitted LogicalIndexMapApply bindings describe views: v3 transposes weight to [768, 1024], then v6 and v7 broadcast the FP32 inputs (v4, v5) across [1024, 1024, 768].

LogicalMul forms the products, and LogicalReduceSum v8 0 sums the trailing K = 768 axis (axis_from_end = 0). The projection is rounded to FP16 in v10 before the final addition. The omitted v12 view broadcasts the bias across [1024, 1024].

Rewriting expands the initial logical program into an e-graph of equivalent expressions discovered under the available rules. An e-class represents one value; its e-nodes are alternative expressions for that value. The view below focuses on four classes from the captured e-graph. Start with the FP32 projection: a sum of products, a matrix multiply, and a transposed product all belong to the same class. Selecting an expression reveals its operand classes, which can themselves contain alternatives.

Extraction chooses compatible alternatives to form a plan; selecting a node in this view only inspects the capture. Equivalence does not imply equal cost: the reduction and matrix forms can lead to different generated kernels or library calls. Retaining alternatives keeps those options open, at the cost of graph growth and plan search (e-graph background). The available rules and search budget limit what can be found; the capture does not guarantee the fastest implementation.

One value, several equivalent expressions

A dashed boundary is an e-class. Each solid card is an e-node inside it. Inspect an expression to see the e-classes it reads.

FP32 projectionLogicalTensor-69[1024, 1024] · FP32

Same value · alternative expressions

3 of 6 captured members shown above. All belong to this e-class.

Sum of products

Reduces the broadcast product over K = 768. The product is another e-class, so its operands can have alternatives too.

Operand e-classes →
  • LogicalTensor-68Broadcast products
  • i64-0Trailing axis (0)
→ LogicalReduceSum
Exact captured node
{
  "function-0-LogicalReduceSum": {
    "op": "LogicalReduceSum",
    "children": [
      "function-1-LogicalMul",
      "primitive-i64-0"
    ],
    "eclass": "LogicalTensor-69",
    "cost": 1.0,
    "subsumed": false
  }
}
Other captured members (3)

These include representation links (logical-of) and additional views. Records retain the original child node IDs.

LogicalIndexMapApplyfunction-14-LogicalIndexMapApply
{
  "function-14-LogicalIndexMapApply": {
    "op": "LogicalIndexMapApply",
    "children": [
      "function-45-logical-of",
      "function-7-IndexMapLit",
      "function-14-shape-of"
    ],
    "eclass": "LogicalTensor-69",
    "cost": 1.0,
    "subsumed": false
  }
}
logical-offunction-13-logical-of
{
  "function-13-logical-of": {
    "op": "logical-of",
    "children": [
      "function-13-LayoutTensorLit"
    ],
    "eclass": "LogicalTensor-69",
    "cost": 1.0,
    "subsumed": false
  }
}
logical-offunction-46-logical-of
{
  "function-46-logical-of": {
    "op": "logical-of",
    "children": [
      "function-100-layout-tensor-list-nth-from-end"
    ],
    "eclass": "LogicalTensor-69",
    "cost": 1.0,
    "subsumed": false
  }
}
From inputs to outputx and weight: [1024, 768] · FP16. Transposed weight is class 52; logical FP32 casts feed the contraction.FP32 projection 69 → FP16 rounding 70 → add bias b[n] → output 87Bias: [1024] · FP16, broadcast across output rows. The rounding step remains before addition for every projection alternative.

Library alternatives. The captured backend retains FP32-storage and FP16-storage cuBLASLt candidates for the contraction. The FP16-storage path requires the research patch described below. These library sites live in additional representation classes.

Selection here only inspects the capture. Extraction later chooses compatible implementations across classes; it does not run every alternative. This view shows four e-classes, not the entire e-graph or every possible program.

Figure 1.4. Selected e-classes from the captured e-graph, with equivalent expressions, operand references, and required FP16 rounding. All 16 members of these four classes are inspectable. Part 2 (forthcoming) follows extraction and measurement of concrete execution plans.

Backend rules attach implementation alternatives such as CublasLtLogicalMatmulSite. In this captured toy example, the FP16-storage Tensor Core alternative requires our backend patch; the unmodified pinned backend handles the explicit-FP32 graph using materialized FP32 operands. Both storage alternatives remain available, and both preserve the FP16 projection cast before bias addition (patch rationale, validation).

What the Semantic IR Makes Available

Figure 1.5 separates how the contraction is expressed (left versus right) from whether equivalent forms are retained together (top versus bottom). Locate TinyGrad, Inductor, and MAX first, then notice that Luminal spans both columns: decomposing MatMul into primitives does not prevent a compiler from recognizing and retaining a matrix operation later.

Two-axis map of the captured contractions. TinyGrad shows primitives; Inductor and MAX show explicit MatMul. Luminal spans both columns with three equivalent expressions—sum of products, direct MatMul, and transposed product—followed by an ellipsis for other captured members.
Figure 1.5. Positions describe the contraction in these captures. The dashed boundary shows three of six captured members of one projection e-class; ≡ means they represent the same value, and … stands for the other three expressions in that class. The formulas use FP32 operands. A complete plan selects compatible members across many e-classes; these are expressions for one value, not complete programs. The earlier e-graph explorer exposes all six members.

The practical question is what information the next optimization needs. Library dispatch benefits from a recognizable contraction; an indexing rewrite needs access to element addresses; choosing among equivalent forms needs a way to compare implementation costs. No single level exposes all of these equally well. Lowering makes some details available while removing structure that another optimization may need.

Retaining alternatives can coexist with progressive lowering: a compiler can preserve several forms at one stage and lower selected forms later. The design choice is where to commit, how much structure to retain, and how much compilation time and memory to spend on that flexibility. Every choice must still respect the numerical contract.

These semantic captures establish what later stages can work with, but do not yet determine the GPU kernel sequence or which design runs fastest. Part 2: Scheduling and Code Generation (forthcoming) follows the same representations into concrete storage choices, library calls, generated kernels, and measured execution.