Warmup Guide — ML Compilers & Intermediate Representations

Zero-to-expert primer for Phase 05. Builds the compiler view of ML from "what is an IR" through ONNX, TVM's Relay/TIR split, and MLIR's dialect system — assuming Phase 04's torch.compile vocabulary.

Table of Contents

Chapter 1: What a Compiler IR Is and Why ML Needs Several

Zero background: a compiler transforms programs between representations. An intermediate representation (IR) is a program format designed for analysis and transformation rather than for humans: explicit dataflow, explicit types, no syntactic sugar. Classic example: LLVM IR sits between C++ and x86, and a hundred optimizations operate on it without caring about either end.

Why ML compilers need multiple IRs: optimization questions live at different altitudes.

  • "Can I fuse conv+BN+ReLU?" is a graph-level question (operators as atoms).
  • "How do I tile this matmul for a 1 MB SRAM?" is a loop-level question (the operator's internals as loops over indices).
  • "Which vector instruction implements this multiply-accumulate?" is instruction-level.

One IR cannot serve all three: an IR that exposes loop indices destroys the operator structure graph passes need, and vice versa. So every serious stack is a stack of IRs with lowering steps between — Relay→TIR in TVM, FX→Inductor-IR→Triton in PyTorch, dialect→dialect in MLIR, ONNX→vendor-IR in every NPU toolchain (Phase 06's DLC is exactly this).

Chapter 2: The Lowering Ladder

Memorize this ladder; every toolchain in the field is an instance of it:

LevelUnit of programExample IRsOptimizations that live here
GraphoperatorsONNX, Relay, FX, MLIR tosa/stablehlofusion, constant folding, layout (NCHW↔NHWC), quantization insertion, partitioning
Tensor/looploops + buffersTIR, Inductor IR, MLIR linalg/affinetiling, vectorization, loop reordering, memory planning, double buffering
TargetinstructionsLLVM IR, PTX, Hexagon ISAregister allocation, instruction selection, scheduling

Lowering is one-way information loss: once a matmul becomes three nested loops, "this was a matmul" is gone (and with it the option to call cuBLAS). Hence the design rule: optimize at the highest level where the optimization is expressible, lower as late as possible. When you see a vendor toolchain do something weird, the explanation is almost always "that information was lost a level above."

Chapter 3: ONNX — The Interchange Layer

What it is: a serialization format (protobuf) for the graph level — a list of nodes, each with an op type from a versioned standard library ("opset"), plus initializers (weights) and typed input/output value-infos. Not a runtime, not a compiler: a contract.

Why it exists: M frameworks × N deployment targets would need M×N converters; with a standard middle format it's M+N. Every NPU toolchain (Qualcomm's included) accepts ONNX precisely because of this economics.

The pieces that bite in practice:

  • Opsets: the op library is versioned; ops change semantics between opsets (e.g. resize/upsample history). Exporter and consumer must agree; "convert to opset 17" is a real and lossy operation.
  • Dynamic axes: dims may be symbolic ("batch"), but many embedded consumers require fully static shapes — you'll re-export with fixed shapes for NPUs.
  • Functions vs primitives: composite ops (GELU) can export as one node (if the opset has it) or a subgraph of primitives — and which one you get changes what the downstream compiler can pattern-match. This is Chapter 2's information-loss principle in action.
  • External data: weights >2 GB overflow protobuf and ship as side files.

Chapter 4: ONNX in Practice — Export, Opsets, and Optimization

Export from PyTorch is itself a capture problem (Phase 04!): the classic torch.onnx.export traces — data-dependent control flow gets baked in; the modern dynamo=True path uses TorchDynamo for honest capture. Either way: always validate — run the ONNX model under onnxruntime on real inputs and compare to PyTorch within tolerance (1e-5 FP32). Silent semantic divergence (training-mode BN, dropout left on, baked branches) is the classic export bug family.

Graph optimization on ONNX (Lab 01's second half): constant folding, identity/dropout elimination, conv-BN fusion, GELU/LayerNorm pattern fusion into single fused ops the runtime has kernels for. onnxruntime applies these at session-load (graph_optimization_ level); onnxsim does shape-inference-driven simplification offline. The lab has you do both and diff the node counts before/after — learn to read what the optimizer did, not trust it blindly.

The accuracy discipline: every transformation step in a deployment pipeline gets a numeric diff gate (PyTorch → ONNX → optimized ONNX → quantized → NPU). When the final number is wrong, the bisection over stages is your debugging procedure — Phase 06 and the capstone build exactly this harness.

Chapter 5: TVM — Relay and TIR

TVM is the canonical open-source instance of the full ladder, which is why it's worth learning even if you deploy with vendor tools.

Relay (graph level): a typed functional IR — every tensor has a shape+dtype in the type system, so shape inference is type checking. Graph passes (fusion, layout transform, quantization) are Relay→Relay rewrites. (TVM's successor IR "Relax" adds first-class dynamic shapes; the concepts transfer.)

TIR (loop level): explicit loops, buffers, and indices. A matmul in TIR is the triple loop you'd write by hand — plus a schedule: a separable description of how to execute those loops (tile, reorder, vectorize, parallelize, bind to GPU blocks/threads).

The compute/schedule separation (TVM's foundational idea, from Halide): what is computed is written once; how is a list of transformations applied to the loop nest. Correctness lives in the compute; performance lives in the schedule; you can search over schedules without risking correctness. This idea — in various syntaxes — is in Inductor, Mojo, and every NPU compiler's tiling engine.

Chapter 6: Scheduling and Auto-Tuning

Why tiling works (the one transformation to deeply understand): a naive matmul streams entire rows/columns through registers, re-reading B from DRAM $O(N)$ times. Tile the loops into blocks sized to fit cache/SRAM and each block of B is loaded once per tile instead of once per element — DRAM traffic drops by the tile factor. On NPUs with software-managed SRAM (no cache to save you), tiling isn't an optimization, it's mandatory correctness — data must be explicitly staged.

The search problem: tile sizes × loop orders × vectorization × unrolling is a space of millions of schedules whose performance is non-convex and hardware-specific. AutoTVM searches it empirically: generate candidates, compile, measure on the real device, fit a cost model, iterate (auto-scheduler/Ansor generates the candidates structurally rather than from hand-written templates). Lab 02 runs this loop and asks you to inspect the winning schedule and explain why it won — measured search beating your intuition is the lesson.

Chapter 7: MLIR — Dialects and Progressive Lowering

The problem MLIR solves: everyone (TVM, XLA, every vendor) was rebuilding the same compiler infrastructure — parsers, pass managers, SSA, verifiers — for their own IR stack. MLIR is infrastructure for building IRs: one framework, many dialects.

A dialect is a namespaced set of ops + types + verification rules. The ones to recognize on sight:

  • stablehlo / tosa — graph-level ML ops (the ONNX-altitude layer)
  • linalg — structured operations on tensors ("this is a matmul" preserved as a property — fusion and tiling are generic over all linalg ops)
  • affine / scf — loops (affine = polyhedral-analyzable loops)
  • memref — buffers with explicit layout (post-bufferization: tensors are SSA values; memrefs are memory)
  • llvm — the exit to LLVM codegen

Progressive lowering: a compilation is a pass pipeline walking dialects downward (stablehlo → linalg → scf → llvm), each step small and verifiable, mixed dialects coexisting in one module mid-flight. This is Chapter 2's ladder made into a reusable library — and it is what Qualcomm, Apple, and Google NPU compilers are actually built on, which is why Lab 03 (IR reading/building via Python bindings) is career-relevant even though you'll rarely write C++ passes yourself.

Reading MLIR (the lab's core skill): SSA values (%0), ops with fully-explicit types (linalg.matmul ins(%A, %B : tensor<128x256xf32>, ...) outs(...)), regions (ops containing blocks of ops — how control flow and loop bodies nest).

Chapter 8: Classic Optimizations Every ML Compiler Runs

Recognize these by their fingerprints in before/after IR dumps:

  • Constant folding: subgraphs with all-constant inputs computed at compile time (weight transposes, shape arithmetic). Fingerprint: nodes vanish, initializers grow.
  • CSE: identical computations deduplicated.
  • DCE: ops whose results are unused removed (often exposes bugs — your "missing op" was dead code).
  • Operator fusion: Phase 04 Chapter 9's memory-traffic argument; at this level done by pattern (vertical: producer→consumer chains; horizontal: same-shaped siblings batched).
  • Layout assignment: NCHW vs NHWC vs blocked layouts (NC4HW4...) chosen per backend; transposes inserted at boundaries — minimizing those transposes is a real optimization problem and a recurring NPU profiling finding (Phase 06).
  • Memory planning: liveness analysis over the graph → buffer reuse plan → peak memory. On NPUs this is again not optional: the plan must fit physical SRAM.

Lab Walkthrough Guidance

Order: Lab 01 (ONNX) → Lab 02 (TVM) → Lab 03 (MLIR) — interchange first, then one full stack, then the infrastructure generalization.

  • Lab 01: export a small transformer; validate numerically before optimizing (the export bug families are the lesson). Then run the optimizer and produce a node-count + op-type diff table; find the conv-BN and GELU fusions in the diff.
  • Lab 02: build the Relay graph, compile unscheduled, measure; then apply manual tiling, measure; then AutoTVM, measure. The three-point comparison (naive/manual/tuned) is the deliverable — keep the numbers.
  • Lab 03: read before you write — dump linalg for a matmul and annotate every line. Then build a small module with the Python bindings and run a pass pipeline (linalg → loops), diffing IR at each stage.

Success Criteria

You are ready for Phase 06 when you can, from memory:

  1. Draw the three-level lowering ladder and place ONNX, Relay, TIR, FX, linalg, and a vendor DLC on it.
  2. State the information-loss principle and give the matmul→loops→"can't call cuBLAS" example.
  3. Explain opsets and two real ONNX export bug families, with the validation discipline that catches them.
  4. Explain compute/schedule separation and why tiling reduces DRAM traffic with the block-reuse argument.
  5. Define an MLIR dialect, name five, and read a linalg.matmul aloud.
  6. Given an IR before/after dump, identify which classic optimization ran.

Interview Q&A

Q: Your ONNX model's accuracy differs from PyTorch by 2%. Walk through the debug. Bisect the pipeline: (1) raw export vs PyTorch on identical inputs — if differing, suspect capture (eval mode? dropout? data-dependent branch baked by trace? opset semantics like resize coordinate modes); (2) optimized vs raw ONNX — a bad fusion or folding (disable optimization levels to isolate); (3) if quantized, that's Phase 03's playbook. Always with fixed seeds and per-layer output diffs, not just end-to-end.

Q: Why do NPU compilers care so much about layout? The MAC array consumes data in a fixed blocked pattern; a mismatched layout means either transpose ops (DRAM round-trips) or strided access (wasted bandwidth). Layout assignment chooses one layout per island and pushes transposes to island boundaries — profiling that shows 30% time in transposes (a real Phase 06 finding) means the assignment failed.

Q: TVM's schedule vs torch.compile's Inductor — compare. Same separation, different defaults: TVM exposes scheduling as user-facing API + offline empirical search per target (minutes-to-hours of tuning, peak results, deploy-time philosophy); Inductor makes scheduling decisions with heuristics + limited autotuning at JIT time (seconds, near-peak, development-loop philosophy). NPU vendor compilers are TVM- shaped: offline, search-heavy, target-specific.

Q: Why did MLIR win over everyone building their own IR stack? Shared infrastructure (SSA, verification, pass management, location tracking) is ~80% of a compiler's engineering and 0% of its differentiation. Dialects let vendors keep their proprietary altitude (a custom NPU dialect) while reusing lowering to/from standard dialects — interop and hiring both improve. The bet: composable dialects beat monolithic IRs, and the ecosystem (stablehlo, IREE, torch-mlir, every new accelerator) confirmed it.

References