Compilers, Graphs and Runtimes

Part 3 ended with a single kernel running near the roofline, which is almost never where a real deployment sits, because a model is not one kernel. It is a few hundred operations wired into a graph, and the thing deciding which kernels run, in what order, over what memory, in what layout, is the compiler.

So: what happens between model.forward() and the instruction stream, and which steps can I influence? My day job at Arm sits on exactly this boundary, ML compilers and libraries for Arm architectures, and the honest summary is that most of the surprising performance in a deployment, good and bad, comes from compiler passes nobody looked at.

Why an ML Compiler Exists at All

The operator times backend explosion

The naive plan for making machine learning fast is to write a good library: someone writes the best convolution anyone has seen, everyone calls it, done. That worked in 2014 and stopped working for a reason you can count.

A modern core operator set is roughly 180 to 200 operations. PyTorch’s Core ATen opset, which torch.export targets after decomposition, sits in that range, and ONNX’s standard opset is similar.The full ATen surface is far larger, well over two thousand operator overloads once you count in-place variants, out variants and dtype specialisations. The point of a "core" opset is to be the small set everything else decomposes into, so a backend has a finite obligation. Now multiply.

The combinatorial obligation of a hand written kernel library, for one plausible product matrix:

At a generous two engineer-days per variant including tests and a benchmark, that is 25,920 engineer-days, about 104 engineer-years. And that is before shapes: a GEMM fast at \(M=1\) is not the code that is fast at \(M=512\).

Libraries respond by cheating, sensibly: cover the twenty operators that matter with enormous care, give everything else a reference implementation. That is why your model runs at 90% of peak until you use one unusual operator, then falls off a cliff. The cliff is the library’s budget being honest with you.

What the compiler actually promises

A compiler attacks the multiplication rather than the cells. Instead of 12,960 kernels you write a decomposition into a few dozen structured primitives, a lowering into loop nests, a set of loop transformations (tile, vectorise, unroll, fuse), and one code generator per instruction set: cost \(O + P + T + B\) instead of \(O \times B \times D \times L\). You pay in quality, because generated code is usually worse than a hand written kernel where someone bothered to write one.Which is why production stacks are hybrids. On Arm a graph compiler does the graph-level work and dispatches GEMM and convolution into Arm Compute Library or similar, generating code itself only for the elementwise and reduction tail.

Takeaway: A kernel library’s cost grows as the product of operators, targets, data types and layouts; a compiler’s cost grows as their sum. That is why every serious deployment stack has a compiler in it, whether or not it calls itself one.

Capturing the Graph

Before a compiler can optimise anything it needs a graph, and your model is not a graph. It is Python.

Tracing, scripting and export

Three ways to turn Python into a graph, and what each one loses:

Tracing is robust because it only ever sees real tensors, and lossy for the same reason: any Python that is not a tensor operation simply vanishes. torch.export instead builds one FX graph in Core ATen IR, plus a signature separating parameters, buffers and user inputs, plus guards recording its assumptions.The critical difference from torch.compile is behaviour on a graph break. torch.compile may give up on a region and run it in eager Python, fine for training throughput and useless on a device with no interpreter. torch.export raises instead, which is annoying and correct. The graph is a flat list of nodes, each with one of six opcodes:

# opcode         name     target                   args
placeholder      x        x                        ()
get_attr         w        linear.weight            ()
call_function    mm       aten.mm.default          (x, w_t)
call_function    add      aten.add.Tensor          (mm, b)
call_function    relu     aten.relu.default        (add,)
output           output   output                   ((relu,),)

Static single assignment, explicit data flow, no hidden state. Everything downstream works on structures of this shape.

ONNX and its sharp edges

ONNX lets a model from one framework run in a runtime written by someone else: a protobuf holding a graph, its initialisers and a set of opset versions. Genuinely useful, with failure modes that will cost you a week if nobody warns you.

The four ONNX problems I hit most often, and what each looks like from the outside:

Opset drift is the format working as designed. LayerNormalization only became a standard ONNX operator at opset 17, so export a transformer at opset 16 and you get a faithful decomposition: ReduceMean, Sub, Pow, ReduceMean, Add, Sqrt, Div, Mul, Add. Nine nodes and eight round trips to memory, while the backend’s fused kernel sits unused because its matcher expects the one-node form.Every serious ONNX consumer therefore ships pattern matchers that reassemble decomposed layer norms, GELUs and attention blocks. They are brittle: swap two commutative operands in the exporter and the pattern stops matching, with no error, only a slowdown. The model is correct, just slow, and the slowness is invisible in the diff.

Shape inference failure quietly destroys partitioning. A Reshape whose shape argument is computed from another tensor cannot be resolved statically, so output dimensions become unknown and propagate that uncertainty through every downstream node. Backends requiring static shapes then decline all of them.

Dynamic control flow

A graph is a directed acyclic graph of tensor operations. A Python if whose condition depends on tensor data is not that, and neither is a while whose trip count depends on a decoded token. Three honest ways out:

1. Specialise. Trace with a fixed condition and ship a graph valid only for that condition, with a guard. Fast, and a correctness landmine if the guard is missing.
2. Represent the control flow. Use higher order operators the IR understands: torch.cond and torch.while_loop, or ONNX If, Loop and Scan. The branches become subgraphs.
3. Split the model. Keep the control flow on the host and compile the branches separately.

Option two is principled and costs something real: most accelerator backends cannot take a subgraph containing an If, so the partitioner cuts there. One data-dependent branch mid-network turns one delegated subgraph into three, and you pay the fallback tax described below.

import torch
from torch import cond

class Gate(torch.nn.Module):
    def forward(self, x):
        # Data-dependent branch: tracing bakes in whichever side ran,
        # and torch.export raises rather than guessing.
        #   if x.sum() > 0: return x * 2
        #   return x - 1
        return cond(x.sum() > 0, lambda t: t * 2, lambda t: t - 1, (x,))

ep = torch.export.export(Gate(), (torch.ones(4),))

The IR Stack and How Lowering Works

Four levels of abstraction

Compilers descend through levels, and at each level “what is a value?” has a different answer. An optimisation belongs at the highest level where it is still expressible.

The four levels every ML compiler passes through, whatever it calls them:

Dialects, StableHLO and Linalg

MLIR made the levels idea practical. Its one big idea is the dialect: a namespaced set of operations, types and attributes that coexists with every other dialect in the same module, so stablehlo.dot_general, linalg.matmul, scf.for and llvm.fmul can all appear in one IR, sometimes in one function mid-conversion.

StableHLO is a portable, versioned operation set derived from XLA’s HLO, with explicit compatibility guarantees. Graph level: whole tensors, no loops. Its value is as a stable contract between frameworks and backends.

Linalg is the structured tensor level. Each operation carries indexing maps (affine maps from the iteration space into each operand) and iterator types (parallel, reduction). That metadata enables generic transformations: a pass can tile, fuse or vectorise any Linalg operation without knowing whether it is a convolution or a layer norm, because the maps say which loops are independent.

Lowering is therefore a sequence of small, locally verifiable rewrites, not one translation. A pass that cannot handle something leaves it alone for a later pass, which is why a half-lowered module shows two dialects side by side.

One lowering step in detail

A linear layer with bias and ReLU, at graph level:

// Graph level: three whole-tensor operations, three result tensors.
%0 = stablehlo.dot_general %x, %w
       : (tensor<1x256xf32>, tensor<256x128xf32>) -> tensor<1x128xf32>
%1 = stablehlo.add %0, %bias_bcast : tensor<1x128xf32>
%2 = stablehlo.maximum %1, %zeros  : tensor<1x128xf32>

Lowering to the structured level does two things at once: it gives the contraction an explicit iteration space, and it lets the two elementwise operations collapse into one, because they now share that space.

// Structured level: one contraction, one fused elementwise producer.
%acc = linalg.fill ins(%c0 : f32) outs(%init : tensor<1x128xf32>)
%mm  = linalg.matmul
         ins(%x, %w : tensor<1x256xf32>, tensor<256x128xf32>)
         outs(%acc  : tensor<1x128xf32>) -> tensor<1x128xf32>
%out = linalg.generic {
         indexing_maps = [affine_map<(m, n) -> (m, n)>,
                          affine_map<(m, n) -> (n)>,
                          affine_map<(m, n) -> (m, n)>],
         iterator_types = ["parallel", "parallel"]}
       ins(%mm, %bias : tensor<1x128xf32>, tensor<128xf32>)
       outs(%init : tensor<1x128xf32>) {
       ^bb0(%a: f32, %b: f32, %o: f32):
         %s = arith.addf %a, %b : f32
         %r = arith.maximumf %s, %c0 : f32
         linalg.yield %r : f32
       } -> tensor<1x128xf32>

Read the indexing_maps: the bias is indexed by n alone, which is precisely the broadcast, expressed as an affine map rather than a materialised tensor. %1 and %2 no longer exist as separate tensors.

The next step tiles linalg.matmul into scf.for loops and fuses %out into the innermost tile, so bias and ReLU apply while the accumulator is still in registers. That is the fusion that saves the memory traffic, and it is only expressible because the earlier step attached iteration spaces to everything.

The Passes That Buy the Speed

Folding, elimination and algebra

Constant folding evaluates anything whose inputs are all constant. The headline case is batch normalisation: at inference it is an affine function with constant parameters, so it folds into the preceding convolution’s weights.

Fifty convolutions in a ResNet means fifty whole-tensor elementwise passes deleted at zero accuracy cost.Not quite zero numerical cost: folding changes operation order, so outputs differ in the last mantissa bits. Harmless in fp32, occasionally not in fp16, where the folded weight overflows if gamma over sigma is large. A good compiler declines the fold if the folded range would saturate.

Dead code elimination removes nodes nobody reads, and exported graphs are full of them: auxiliary training heads, unused masks, the second element of a tuple you never indexed. Common subexpression elimination merges identical computations on identical inputs, which matters because attention implementations routinely build the same positional bias twice, once per call site.

Algebraic simplification rewrites into cheaper equivalents: x * 1 and x + 0 disappear, transpose(transpose(x)) becomes x, two stacked reshapes become one, division by a constant becomes multiplication by its reciprocal. These run repeatedly, because each other pass creates fresh debris for them to clear.

Fusion, vertical and horizontal

Fusion is the highest-value graph pass, and the reason is bandwidth, not arithmetic.

Vertical fusion merges a producer into its consumer. Take conv -> bias -> relu with a \(1 \times 64 \times 56 \times 56\) fp32 activation: 200,704 elements, 802,816 bytes, call it 0.80 MB.

• Unfused: the convolution writes 0.80 MB; the bias add reads 0.80 and writes 0.80; the ReLU reads 0.80 and writes 0.80. Beyond the convolution’s own unavoidable write, that is 3.21 MB of extra traffic.
• Fused: bias and ReLU apply to the accumulator before it leaves registers. Zero extra traffic.

On an illustrative device with 20 GB/s of bandwidth, 3.21 MB costs \(3.21 \times 10^6 / (20 \times 10^9) = 161\ \mu\text{s}\). The convolution is \(3 \times 3 \times 64 \times 64 \times 56 \times 56 = 115.6\) million MACs, 231 MFLOP, which at an assumed 200 GFLOP/s takes 1.16 ms. Fusion removes 14% of that layer’s time, fifty times over.

The mechanism is arithmetic intensity, from Part 1: a standalone ReLU reads 4 bytes, writes 4 bytes and does one comparison, 0.125 FLOP/byte, memory bound on any hardware ever built. Fusing it into a producer does not make it faster, it makes it free.

Horizontal fusion merges siblings that share an input. Attention’s three projections are the canonical case: instead of three GEMMs of \([128 \times 768] \times [768 \times 768]\) each re-reading the same 393 KB activation, concatenate the weights into \([768 \times 2304]\) and run one GEMM, reading the activation once, saving 0.79 MB, and paying one launch instead of three.

Layout, propagation and scheduling

Every backend has a preferred layout. Arm NEON and SVE want NHWC for convolution, because the channel dimension is then contiguous and vectorises cleanly, and int8 dot product instructions want four channels adjacent. Some GPU paths want NCHW; blocked layouts like NC8HW8 interleave a channel block to match a vector width exactly.

When two adjacent operations disagree the compiler inserts a transpose, which is pure overhead: the 0.80 MB tensor above costs a read and a write, 1.61 MB, 80 µs at 20 GB/s. Once per layer across 50 layers is 4 ms of doing nothing.

Layout assignment is therefore global. The compiler labels every node so total transpose cost is minimised, then propagates labels so runs of compatible operations share a layout and conversions get pushed to the graph boundary, ideally into preprocessing where they are free. This is the pass I have seen produce the largest unexplained regressions: one operator with only an NCHW implementation forces two transposes per layer around a whole region.

Scheduling picks execution order, which for a branching graph changes which tensors are simultaneously live and whether a weight prefetch overlaps compute. It and memory planning are one problem seen from two sides, so a new scheduling heuristic can move peak memory by 40% without changing a kernel.

Takeaway: Fusion, layout and scheduling have the largest effect of any graph passes, and all three are about memory traffic, not arithmetic. A pass that is not reducing bytes moved or kernel launches will probably not save you.

Memory Planning and Why Peak Is What Kills You

Liveness analysis over the graph

On a server you allocate from a caching allocator and stop thinking. On a part with 8 MiB of usable SRAM there is no virtual memory, no swap and no out-of-memory killer: the allocation either fits at initialisation or the product does not work. What matters is peak memory, not total. Fix a topological order; a tensor is then live from when its producer starts writing to when its last consumer finishes reading, and two tensors may share bytes exactly when their intervals do not overlap.

Tensors in an illustrative seven-operation backbone slice with a long skip, int8 activations, batch 1:

Give every tensor its own buffer and you need 13,056 KiB, 12.75 MiB, which does not fit.

Live bytes at each step of the execution order, the sum that peak memory actually measures:

Peak is 8192 KiB, exactly 8 MiB. Liveness analysis has recovered 4.75 MiB without touching a kernel.

Greedy allocation into an arena

Knowing the peak is not achieving it: you must assign each tensor a byte offset in one arena so tensors overlapping in time never overlap in space. That is the dynamic storage allocation problem, equivalent to two-dimensional strip packing and NP-hard, so planners go greedy: sort by size descending, place each at the lowest offset colliding with no already-placed tensor whose interval overlaps.TFLite's arena planner, ExecuTorch's memory planning passes and TVM's storage rewrite all use variants of this. Size descending is usual because large tensors are what fragment the arena, so placing them first keeps the holes fillable by small ones.

Greedy offset assignment, largest tensor placed first, for the tensors above:

t2 sits at offset 0 because it dies at op4 and t6 is not born until op6; t7 reuses t1’s bytes for the same reason. The arena is 8192 KiB, equal to the peak, so this plan has zero fragmentation. That is common but not guaranteed: greedy plans routinely land 5 to 15% above the peak, which is always a lower bound.

Eight megabytes and what to do when you miss

The plan needs exactly 8 MiB and the device has 8 MiB, so it does not fit: weights, runtime state and DMA scratch also have to live somewhere. You have missed by a little and the product does not ship. Ordinary Tuesday.

The largest line is t6, and a channel-axis concatenation need not copy anything. If the planner places t1 and t5 adjacently in the right order, t6 is a view over those same 4096 KiB and op6 becomes a no-op. This is concat elision, and every mature edge compiler does it. Op6’s live set is then 4096 KiB and op7’s is 5120: new peak 5 MiB, three megabytes recovered by deleting an operation rather than optimising one.

The escape ladder, in the order I try it:

1. Elide copies: concat, split, slice and reshape can often be views.
2. Reorder the schedule so large tensors are born later and die sooner.
3. Recompute rather than store, when recomputation is cheap against the bandwidth saved.
4. Tile the graph: push a spatial strip through several layers at once, so intermediates are tiles rather than whole feature maps. This is what fits a 512x512 input in a few hundred KB.
5. Quantise the activations, which quarters everything at once (Part 5).
6. Spill the largest buffers to external DRAM, at perhaps a 10x bandwidth penalty.

Takeaway: Peak memory, not total, decides whether a model fits. Liveness plus greedy arena allocation typically recovers 2 to 3x over naive allocation, and deleting copies, especially concatenations, is usually worth more than any allocator cleverness.

Autotuning, Cost Models and the Compile Budget

Once the graph is settled, each kernel still needs a schedule: tile sizes, loop order, unroll factors, vector widths, whether to stage a tile in shared memory. AutoTVM’s model was a hand written template with knobs, searched by measuring on the real device.

Count the space for a \(1024^3\) GEMM with M, N and K each split into three levels. Since \(1024 = 2^{10}\), the ordered factor triples per axis number \(\binom{12}{2} = 66\), so three axes give \(66^3 = 287{,}496\). Times six unroll settings and two vectorisation settings: about 3.4 million configurations for one operator at one shape. At 100 ms per trial, realistic once you include compilation, transfer and timed repetitions, exhaustive search is 344,995 s, or 96 hours, on one target.

So nobody searches exhaustively. The cost model is a learned regressor over schedule features: loop extents, estimated cache line accesses, vector width, arithmetic intensity, unrolled instruction count. It does not need accurate runtimes, only correct ranking, since its job is to choose the next batch to measure, so AutoTVM measures 500 to 2,000 configurations instead of millions.

Ansor removed the templates. It generates sketches from the mathematical definition of the computation, annotates them randomly, and searches by evolutionary mutation under a learned cost model. It also adds a task scheduler spending a global budget across subgraphs by estimated end-to-end gain, which is the right framing: tuning an operator worth 2% of runtime to perfection wastes the budget.

Tuning budget for one network on one target, at an assumed 0.3 s amortised per measured trial:

Three hours fits overnight; twenty hours does not fit a per-commit CI job. Hence the central operational fact: autotuning is a build artefact, not a build step. Tune once, cache.

The cache is a tuning log keyed by operator, full shape and dtype signature, and a target string such as llvm -mtriple=aarch64-linux-gnu -mattr=+neon,+dotprod. A hit gives the tuned schedule; a miss silently gives a fallback, commonly 2 to 10x slower. Here is the trap I have walked into personally. Add +sve because you are exploring a new part. Change a sequence length from 128 to 127. Bump the compiler and have it normalise the target string differently. Each is a total cache miss, each produces a correct model, and each appears as an overnight latency regression with no code change to blame.The defence is cheap: make a cache miss loud. Fail the build, or at minimum warn per missed operator with the key it looked for, and assert the hit rate in CI. A silent fallback is the worst default: the failure is invisible in every signal except the benchmark.

Takeaway: Autotuning is search under a wall-clock budget, so cost models exist to avoid measurement and caches exist to avoid search. Treat a tuned schedule cache as a versioned build artefact, and make misses fail loudly rather than degrade quietly.

Runtimes, Delegates and Subgraph Partitioning

A compiled graph still needs something to execute it, allocate its arena and talk to the driver. Every runtime has the same core abstraction under a different name.

The same idea, five vocabularies:

The mechanism is subgraph partitioning. The runtime asks each backend, in priority order, which nodes it can execute. Support is per-attribute, not per-operator-name: a backend refuses a node because the dilation is unusual, the dtype is fp16 rather than int8, or an axis is dynamic. It then forms maximal connected subgraphs of supported nodes, replaces each with one “call this backend” node, and leaves the rest to the CPU. Most partitioners also enforce a minimum subgraph size, since delegating two isolated nodes costs more in handoff than it saves.

The fallback tax

Take a 120 node graph, 108 supported, 8 ms of compute. Each delegated partition is entered and exited once, and each of those boundaries costs a synchronisation plus, usually, a copy of the boundary tensor out of the backend’s memory space and back. For a 1 MiB tensor at an assumed 8 GB/s one copy is 131 µs, so a round trip plus roughly 50 µs of synchronisation is about 0.31 ms per boundary.

The same 108 supported nodes, arranged two ways:

Identical model, identical hardware, identical supported-operator set, 1.9x difference in latency decided by where the unsupported nodes sit, and the second case is worse than the table shows because fusion cannot cross a boundary either. The pathological version is one unsupported node inside a residual block, splitting the network in half.

I have spent more time chasing single unsupported operators than almost any other category of performance bug, and the fix is nearly always: implement it in the backend, rewrite the model to use supported operators, or accept CPU for the whole thing. Read the partitioning report first. TFLite logs how many nodes were delegated and in how many partitions, ONNX Runtime dumps node placement per execution provider, and ExecuTorch prints a partition summary. One line saying “13 partitions” is worth a day of profiling.

Ahead of time, just in time, and the version trap

Ahead of time against just in time, for an edge target:

AOT is the only option when the device has no room for a compiler, which covers most of embedded: a microcontroller runs a pre-planned arena and a fixed kernel list, and the memory plan is inspectable at build time. The practical middle ground is AOT for the graph plus a warm-up and an on-disk cache for anything JIT, so the first real request is never the first compile.

Serialisation makes this an operations problem. A serialised TensorRT engine is tied to the TensorRT version, the CUDA version and the compute capability it was built for, often the SKU. Move it and deserialisation fails, or silently rebuilds and your measured cold start becomes two minutes. The mobile equivalent is a delegate whose behaviour depends on a vendor driver that changes with an OTA update you do not control.

Pin all five of these, and record them in the artefact manifest:

CI should assert both a latency bound and an output hash against a golden. Without the hash you discover, three weeks late, that a minor version bump changed a fusion decision, which changed the accumulation order, which moved an accuracy metric by half a point.Bit-exactness across compiler versions is not reasonable to demand, since fusion legitimately changes accumulation order. What is reasonable is a tolerance plus an accuracy check on a held-out set, run on every dependency bump: the goal is not to prevent change but to notice it.

Takeaway: The runtime gives each backend the subgraphs it claims and runs the rest on CPU, so latency depends less on how many operators are supported than on how the unsupported ones are distributed. Check the partition count before you profile anything.

When the Compiler Makes It Slower

This happens often enough that I keep a checklist. Work it in order; the early items catch most of it.

1. Check the measurement. Same shapes, same threads, same CPU governor, enough iterations, first run discarded.
2. Confirm it compiled at all. Stacks fall back to a reference path silently, and zero delegated nodes is the most common cause of all.
3. Count the partitions. One is good, thirteen is a disaster. Find the unsupported operator behind each split.
4. Diff the final graph against the input. Look for what the compiler added: transposes, quantise and dequantise pairs, copies.
5. Check the layout. An NCHW pipeline feeding an NHWC backend costs a transpose per inference or, worse, per layer.
6. Check the schedule cache. Confirm the target string and shape signature match a tuning log entry exactly.
7. Check for dynamic shapes. One dynamic dimension disables tiling, vectorisation and static memory planning at once; compile a fixed shape as a control.
8. Compare arena size to your memory hierarchy. Past SRAM or a cache level, you traded arithmetic for bandwidth and lost.
9. Check for numerics-driven refusals. A strict-precision or deterministic flag disables fast paths and fusions, quietly.
10. Profile the kernels, not the graph, since graph totals hide the one kernel taking 80% of the time: Part 3 for the kernel, Part 10 for the method.
11. Only then file a bug, with the smallest graph that reproduces it. A two-node reproducer gets fixed; a whole model does not.

A Few Problems to Work

Problem 1. An inception-style block takes an input x of 1024 KiB, feeds it to four sequential branches each producing 512 KiB, then concatenates the four into a 2048 KiB result. x is live until the last branch consumes it. Compute peak live memory with a materialised concatenation, then with concat elision, and give the percentage saving.

Problem 2. A 200 node graph runs in 8 ms fully delegated. Your backend supports 180 nodes and each partition boundary costs 0.31 ms. First, compute latency if the 20 unsupported nodes are contiguous, and if they scatter into 20 delegated partitions. Second, if the CPU-only baseline is 22 ms, at how many delegated partitions does delegation stop paying?

Problem 3. You are tuning a network with 40 distinct operator shapes across 5 target configurations, at 0.3 s per measured trial and 1,200 guided trials per task. Does the tuning run fit a 12 hour overnight window, and if a target string change invalidates the cache for 30% of tasks on every target, what does each incident cost?

Problem 4. A 50 layer network runs on a backend that wants NHWC. Twelve layers have only NCHW kernels and are scattered, never adjacent. Each activation is 0.80 MB and bandwidth is 20 GB/s. How many transposes does layout assignment insert, what do they cost, and what changes if the twelve are contiguous instead?

What’s Next

The compiler decides which kernels run, in what order, over which bytes. Almost every lever in this chapter, fusion, layout, memory planning, partitioning, is really a lever on bytes moved. The remaining way to move far fewer bytes is to make each one carry more information, by changing the numeric format itself. On most edge deployments that buys more than any other single change: an int8 model moves a quarter of the bytes of an fp32 one and runs on hardware paths that do not exist for floats at all.

That’s all for Part 4! For Part 5, on quantisation, click here.