Phase 05 — ML Compilers & IRs
Difficulty: ⭐⭐⭐⭐⭐
Estimated Time: 3 weeks (70–90 hours)
Roles Supported: Framework & compiler engineers — required for writing NPU backend compilers, ONNX export pipelines, and cross-framework optimization toolchains
Why This Phase Exists
The ML compiler stack is the path from a PyTorch model to a running binary on a Qualcomm NPU. The chain is:
PyTorch Model
→ torch.export / torch.onnx.export
→ ONNX / MLIR / StableHLO (portable IRs)
→ SNPE/QNN compiler (Qualcomm-specific)
→ .dlc or .bin (NPU binary)
→ Hexagon runtime
Understanding each IR is not optional if you want to debug why a model produces wrong results after compilation, or why 3 ops are missing from the NPU coverage report, or why the compiled model is slower than expected despite the op coverage being 95%.
This phase builds working knowledge of ONNX, TVM, and MLIR — the three dominant ML compiler IRs in production systems. ONNX is what most pipelines export to. TVM is a complete compiler that demonstrates autotuning and schedule optimization. MLIR is the foundation of many new compilers including IREE (used for Qualcomm AI Engine).
Concepts
- ONNX (Open Neural Network Exchange): protobuf-based IR; nodes are
NodeProtowith op_type, inputs, outputs, attributes; opset versioning (opset 17 = standard); computation graph is a topological list of nodes - ONNX opset: each op version specifies exact semantics; mismatches between export opset and runtime opset cause failures
onnxoptimizer: graph-level rewrites for ONNX graphs (constant folding, reshape fusion, conv+BN fusion); similar to FX passes but for ONNXonnxruntime: Microsoft's ONNX execution engine; multiple execution providers (CPU, CUDA, CoreML, QNN)- TVM: complete ML compiler stack; Relay (high-level IR) → TIR (tensor IR) → codegen → machine code; includes AutoTVM (search-based schedule tuning) and Ansor (auto-scheduler)
- Relay IR: functional, typed, higher-order; supports data types, shapes, control flow; analogous to FX but more formal
- TIR (Tensor IR): imperative, loop-based; describes memory access patterns, loop tiling, vectorization; where most performance optimization happens
- Schedule in TVM: explicit loop transformation specification (tile, vectorize, parallelize, reorder, unroll); AutoTVM/Ansor searches the schedule space automatically
- MLIR (Multi-Level IR): modular IR framework; dialects define ops and types; transformation passes lower between dialects; core of LLVM 17+ and many ML compilers
- MLIR dialects:
linalg(named mathematical operations),affine(affine maps on memrefs),vector(SIMD),memref(memory regions),arith(scalar arithmetic),func(functions),tosa(Tensor Operator Set Architecture — used by TFLite and Qualcomm) - TOSA (Tensor Operator Set Architecture): small, well-defined op set targeting hardware accelerators; used by TFLite converter, IREE, and Qualcomm's pipeline
- IREE (Intermediate Representation Execution Environment): end-to-end ML compiler; ingests MLIR/StableHLO, outputs code for CPU/GPU/NPU; used for Qualcomm AI Engine deployments
- StableHLO: XLA's HLO dialect stabilized as a cross-framework standard; JAX/TF/PyTorch can all emit StableHLO
Labs
Lab 01 — ONNX Export Pipeline + Optimization Passes
| Field | Value |
|---|---|
| Goal | Build a complete ONNX export, validation, and optimization pipeline for LLaMA-style models; implement 3 custom graph rewrite passes; validate accuracy end-to-end |
| Concepts | torch.onnx.export, onnx.checker, onnxoptimizer, onnxruntime, opset compatibility, dynamic axes, custom ONNX operators |
| Steps | 1. Export LLaMA-tiny to ONNX with dynamic axes for batch and sequence; 2. Run onnx.checker.check_model() and fix any violations; 3. Implement 3 custom onnxoptimizer passes: (a) fuse consecutive Reshape ops, (b) eliminate identity ops, (c) fold constant Gather nodes; 4. Validate accuracy: ONNX Runtime output vs PyTorch within 1e-4 on 100 inputs; 5. Compare latency: PyTorch eager vs ONNX Runtime CPU vs ONNX Runtime CUDA |
| Stack | PyTorch 2.3+, onnx, onnxoptimizer, onnxruntime, onnxruntime-gpu |
| Datasets | Synthetic + WikiText-2 val (first 100 sequences) |
| Output | llama_tiny.onnx; optimized llama_tiny_opt.onnx; accuracy/latency comparison table; visualization of graph before/after optimization |
| How to Test | pytest test_lab.py — checks: ONNX model passes checker, output matches within 1e-4, optimized model has fewer nodes, dynamic axes work for different batch/seq sizes |
| Talking Points | What are dynamic axes in ONNX? Why does a Reshape → Reshape sequence arise from PyTorch export? How do custom ONNX ops work when the runtime doesn't support them? |
| Resume Bullet | Built ONNX export + optimization pipeline for LLaMA-style model; 3 custom graph passes reduced node count by 18%; ONNX Runtime achieved 1.3x latency improvement over PyTorch eager on CPU |
| Extensions | Add QDQ nodes for INT8 ONNX quantization; add TensorRT EP execution; integrate with ONNX Runtime's QNN Execution Provider |
Lab 02 — TVM Compilation Pipeline
| Field | Value |
|---|---|
| Goal | Compile a ResNet-50 and BERT-base through the full TVM pipeline (Relay → TIR → AutoTVM tuning → deploy), demonstrating 2x+ speedup over PyTorch eager via schedule optimization |
| Concepts | TVM Relay IR, relay.frontend.from_pytorch, relay.build, tuning records, autotvm.LocalRunner, ansor.TaskScheduler, TVM runtime module, cross-compilation for ARM |
| Steps | 1. Import PyTorch model to Relay IR via relay.frontend.from_pytorch; 2. Run shape inference and type checking; 3. Apply standard Relay optimizations (FuseOps, FoldConstant, EliminateCommonSubexpr); 4. Build without tuning (baseline); 5. Extract tuning tasks and run AutoTVM with 100 trials; 6. Build with tuning records; 7. Compare: PyTorch eager vs TVM untuned vs TVM tuned |
| Stack | TVM 0.15+, PyTorch 2.3+, torchvision |
| Datasets | ImageNet val subset (timing) |
| Output | TVM compiled modules for ResNet-50 and BERT-base; timing table showing speedup; tuning logs; cross-compilation target for llvm -mcpu=cortex-a55 (simulating ARM edge device) |
| How to Test | pytest test_lab.py — checks: Relay model passes type checker, output matches PyTorch within 1e-4, tuned model faster than untuned, cross-compile succeeds without actual ARM device |
| Talking Points | What is a schedule in TVM? How does AutoTVM search the schedule space? What is the Relay FuseOps pass and how does it decide what to fuse? |
| Resume Bullet | Compiled ResNet-50 through TVM with AutoTVM tuning; achieved 2.1x throughput improvement over PyTorch eager on x86 CPU; demonstrated ARM cross-compilation for edge deployment |
| Extensions | Add Ansor (auto-scheduler) comparison; enable CUDA VTA (virtual tensor accelerator) target; implement custom TIR pass for loop reordering |
Lab 03 — MLIR Dialect and Lowering Pass
| Field | Value |
|---|---|
| Goal | Write a simple MLIR transformation pass using Python bindings that lowers a subset of linalg ops to affine dialect; then lower to LLVM IR and compile to native code |
| Concepts | MLIR Python bindings (mlir-python-bindings), dialects, IR builder API, pass pipeline, linalg.generic, affine.for, memref.alloc, arith ops, bufferization |
| Steps | 1. Build a simple 3-op MLIR program (matmul + bias + relu) using Python bindings; 2. Print the linalg dialect representation; 3. Apply the standard lowering pipeline: linalg-bufferize → convert-linalg-to-affine → lower-affine → convert-to-llvm; 4. Write a custom pass: AffineLoopUnroll that unrolls innermost loops by factor 4; 5. JIT compile and execute via mlir.execution_engine.ExecutionEngine; 6. Verify output matches NumPy reference |
| Stack | mlir-python-bindings (from LLVM 17+), llvmlite |
| Datasets | Synthetic matrices |
| Output | Complete lowering pipeline from linalg → affine → llvm → native; custom unroll pass; JIT execution with correctness verification |
| How to Test | pytest test_lab.py — checks: linalg IR is valid, lowered affine IR has expected loop structure, JIT output matches NumPy within 1e-5, custom unroll reduces loop iteration count by 4× |
| Talking Points | What is a dialect in MLIR? How does progressive lowering work? Why does linalg.generic use an indexing_map? How does MLIR relate to Qualcomm's IREE/QNN compiler? |
| Resume Bullet | Built MLIR transformation pipeline with custom affine loop unroll pass; demonstrated linalg→affine→LLVM lowering for matmul+bias+relu, producing JIT-compiled code matching NumPy reference within 1e-5 |
| Extensions | Target TOSA dialect (for Qualcomm pipeline); write a SPIR-V lowering pass; implement vectorization pass using vector dialect |
Deliverables Checklist
- ONNX pipeline with 3 custom passes and accuracy/latency table
- TVM compilation with AutoTVM tuning showing ≥2× speedup
- MLIR lowering pipeline with custom pass and JIT execution
-
All
test_lab.pysuites pass
Interview Relevance
- "What is the difference between ONNX and MLIR?" — you can explain: ONNX = serializable exchange format, MLIR = compilation framework with progressive lowering; they solve different problems
- "How would you add a new op to the QNN compiler?" — you know to: register in ONNX (for export), implement in TOSA dialect (MLIR), register as SNPE custom op (C++)
- "Why does a model run correctly in PyTorch but produce NaN in TVM?" — you know to check: type casting (FP32 vs FP16 accumulation), padding behavior differences, reduction order numerics