The Hitchhiker's Guide — Phase 05: ML Compilers & Intermediate Representations

"ONNX is the lingua franca of ML models. It is also the language everyone speaks but nobody reads." — ML compiler team

Section 1: Why ML Compilers Exist

A PyTorch model is a Python program that calls CUDA kernels. This is flexible but suboptimal:

1. Each op is a separate CUDA kernel launch (overhead: ~5 µs per launch)
2. Elementwise op sequences allocate intermediate tensors (memory overhead)
3. The scheduler doesn't know about future ops — can't make global decisions

ML compilers solve this by:

1. Capturing the full computation graph (no Python interpreter overhead)
2. Optimizing globally across ops (operator fusion, constant folding, layout optimization)
3. Generating hardware-specific code (Triton, CUDA C, HTP microcode, LLVM IR)

The result: compiled models run 1.5–10× faster than eager mode, depending on how memory-bound the workload is.

Section 2: ONNX — Open Neural Network Exchange

What ONNX Is

ONNX is a graph-based IR (intermediate representation) with:

• Nodes: operations (Conv, Gemm, LayerNormalization, etc.)
• Inputs/Outputs: typed tensors with optional shape information
• Initializers: constant weights (stored in the proto)
• Opset: versioned operation set (opset 17 = ONNX spec version 17)

ONNX is a serialization format — it stores the graph as a .onnx file (protobuf binary) that any ONNX-compatible runtime can load.

Exporting from PyTorch

import torch
import torch.onnx

model = MyModel()
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

# Old API (still common)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}  # dynamic batch
)

# New API (torch.onnx.dynamo_export — uses torch.compile under the hood)
from torch.onnx import dynamo_export
export_output = dynamo_export(
    model,
    dummy_input,
    export_options=torch.onnx.ExportOptions(dynamic_shapes=True)
)
export_output.save("model_dynamo.onnx")

Inspecting ONNX Graphs

import onnx
from onnx import numpy_helper

model = onnx.load("model.onnx")
onnx.checker.check_model(model)  # validate the graph

# Inspect nodes
for node in model.graph.node:
    print(f"Op: {node.op_type}, inputs: {list(node.input)}, outputs: {list(node.output)}")

# Inspect tensors
for init in model.graph.initializer:
    arr = numpy_helper.to_array(init)
    print(f"Weight: {init.name}, shape: {arr.shape}, dtype: {arr.dtype}")

# Use netron for visual inspection:
# pip install netron && netron model.onnx

ONNX Graph Optimization

from onnxoptimizer import optimize

# Standard optimization passes
optimized = optimize(model, [
    'eliminate_unused_initializer',
    'eliminate_nop_transpose',
    'eliminate_nop_pad',
    'fuse_bn_into_conv',          # batchnorm → folded into conv weights
    'fuse_consecutive_transposes',
    'fuse_matmul_add_bias_into_gemm',
    'fuse_pad_into_conv',
    'eliminate_deadend',
])

# onnxsim (onnx-simplifier) — more aggressive
import onnxsim
simplified, check = onnxsim.simplify(model)
assert check, "Simplified model is invalid"

Pitfall: Not all ONNX opset versions are supported by all runtimes. QNN supports a subset of opset 17. Always validate with the target runtime after optimization.

ONNX Runtime

import onnxruntime as ort
import numpy as np

# CPU inference
session = ort.InferenceSession("model.onnx", providers=['CPUExecutionProvider'])

# CUDA inference
session = ort.InferenceSession("model.onnx", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

# Run inference
input_name = session.get_inputs()[0].name
output = session.run(None, {input_name: np.random.randn(1, 3, 224, 224).astype(np.float32)})

Section 3: TVM — Tensor Virtual Machine

TVM is an open-source ML compiler that lowers a model graph → Relay IR → TE (Tensor Expressions) → TIR (Tensor IR) → backend code (LLVM, CUDA, Metal, HTP).

TVM Compilation Stack

ONNX / TorchScript / TFLite
          ↓
    Relay IR (high-level graph IR)
          ↓
    [Relay passes: fold, inline, quantize]
          ↓
    TE (Tensor Expressions) per operator
          ↓
    [AutoTVM / AutoScheduler — search for best schedule]
          ↓
    TIR (Tensor IR — loop nest representation)
          ↓
    [TIR passes: loop fusion, vectorization, unrolling]
          ↓
    Backend code (LLVM, NVPTX, Metal, HEX)

Basic TVM Example

import tvm
from tvm import relay
import tvm.relay.testing
import onnx

# Load ONNX model into Relay
onnx_model = onnx.load("model.onnx")
shape_dict = {"input": (1, 3, 224, 224)}
mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)

# Apply Relay optimization passes
seq = tvm.transform.Sequential([
    relay.transform.InferType(),
    relay.transform.FoldConstant(),
    relay.transform.FastMath(),
    relay.transform.FoldScaleAxis(),
    relay.transform.CanonicalizeOps(),
    relay.transform.AlterOpLayout(),  # layout optimization (NHWC→NCHW, etc.)
])
with tvm.transform.PassContext(opt_level=3):
    mod = seq(mod)

# Compile for target
target = tvm.target.cuda()  # or "llvm" for CPU, "hexagon" for Qualcomm DSP
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target=target, params=params)

# Deploy
dev = tvm.device(str(target), 0)
m = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
m.set_input("input", tvm.nd.array(np.random.randn(1, 3, 224, 224).astype("float32")))
m.run()
output = m.get_output(0).numpy()

TVM AutoScheduler (Ansor)

AutoScheduler searches for optimal tile sizes, unroll factors, vectorization decisions automatically — no hand-written schedules needed:

from tvm import auto_scheduler

log_file = "resnet50.json"

# Measure on hardware to find best schedule
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
    num_measure_trials=200,  # per task
    runner=auto_scheduler.LocalRunner(repeat=3, min_repeat_ms=150),
    measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
tuner.tune(tune_option)

# Compile with tuned schedule
with auto_scheduler.ApplyHistoryBest(log_file):
    with tvm.transform.PassContext(opt_level=3, config={"relay.backend.use_auto_scheduler": True}):
        lib = relay.build(mod, target=target, params=params)

Section 4: MLIR — Multi-Level Intermediate Representation

MLIR is a compiler infrastructure (from LLVM project) that provides:

• Dialects: modular IR namespaces (e.g., linalg, arith, memref, gpu, tosa)
• Transformations: composable passes that lower between dialects
• Extensibility: you can define a new dialect for your hardware

PyTorch uses MLIR internally (via torch-mlir and iree-compiler) for export to mobile and edge targets.

The Dialect Lowering Hierarchy

torch dialect         (PyTorch ops: torch.aten.linear, torch.aten.relu)
      ↓ (torch-mlir)
tosa dialect          (Tensor Operator Set Architecture — hardware-agnostic)
      ↓
linalg dialect        (structured loops: linalg.matmul, linalg.generic)
      ↓
arith + memref        (arithmetic + memory buffer ops)
      ↓
llvm dialect          (LLVM IR)
      ↓
Target ISA (x86, ARM, HEX)

MLIR Python Bindings

from mlir.ir import Context, Module, InsertionPoint
from mlir.dialects import arith, func

with Context() as ctx:
    module = Module.create()
    with InsertionPoint(module.body):
        # Define a function: (i32, i32) -> i32
        @func.FuncOp.from_py_func(
            arith.IntegerType.get_signless(32),
            arith.IntegerType.get_signless(32)
        )
        def add(a, b):
            result = arith.AddIOp(a, b)
            return result.result

print(module)
# prints:
# module {
#   func.func @add(%arg0: i32, %arg1: i32) -> i32 {
#     %0 = arith.addi %arg0, %arg1 : i32
#     return %0 : i32
#   }
# }

torch-mlir: PyTorch → MLIR

import torch
import torch_mlir

model = MyModel()
model.eval()

# Export to MLIR (TOSA dialect)
module = torch_mlir.compile(
    model,
    torch.randn(1, 3, 224, 224),
    output_type="tosa",
)
print(module.operation.get_asm())

# Export to MLIR (linalg-on-tensors)
module = torch_mlir.compile(
    model,
    torch.randn(1, 3, 224, 224),
    output_type="linalg-on-tensors",
)

Section 5: Qualcomm-Specific Compiler Path

Qualcomm's compilation path for Snapdragon:

PyTorch model
    ↓ (torch.onnx.export or torch.export)
ONNX graph
    ↓ (qai-hub-models or SNPE converter)
QNN DLC (Deep Learning Container)
    ↓ (QNN runtime on device)
HTP (Hexagon Tensor Processor) execution

Key Qualcomm tools:

• snpe-onnx-to-dlc: converts ONNX → SNPE DLC (older, SNPE 2.x)
• qnn-onnx-converter: converts ONNX → QNN model (newer, QNN SDK)
• qai_hub.upload_model(): upload to AI Hub for cloud compilation + profiling
• qai_hub.compile_model(): compile for specific device (e.g., "Snapdragon 8 Gen 3")

Qualcomm-specific constraints:

1. Op support: not all ONNX ops are supported on HTP; check QNN op support matrix
2. Data types: HTP prefers INT8/INT16 weights; FP16 ops run on HVX or Cortex-A cores
3. Static shapes: HTP compilation requires static input shapes (no dynamic batch)
4. Graph partitioning: ops not supported on HTP fall back to CPU; too many fallbacks → poor performance

Section 6: Interview Pitfalls

Section 7: Resources

1. ONNX specification — https://onnx.ai/onnx/intro/concepts.html
2. onnxoptimizer passes — https://github.com/onnx/optimizer/tree/master/onnxoptimizer/passes
3. TVM tutorial — https://tvm.apache.org/docs/tutorial/index.html
4. AutoScheduler (Ansor) paper — Lianmin Zheng et al., "Ansor: Generating High-Performance Tensor Programs for Deep Learning" (OSDI 2020)
5. MLIR primer — https://mlir.llvm.org/docs/Tutorials/Toy/
6. torch-mlir — https://github.com/llvm/torch-mlir
7. QNN SDK documentation — https://docs.qualcomm.com/bundle/publicresource/topics/80-63442-2/
8. "The Deep Learning Compiler: A Comprehensive Survey" — Li et al., 2021 (arXiv:2002.03794)