Lab 02 — TVM Compilation
Phase: 05 — ML Compilers & IR | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours
Apache TVM compiles ML models to optimized native code for CPU, GPU, and custom accelerators.
What you build
- Import PyTorch model into TVM Relay IR
- Apply TVM optimization passes (fusion, layout transform, quantization)
- Compile to native code with
tvm.build - Benchmark TVM vs PyTorch on target hardware
- Analyze Relay IR for supported vs unsupported ops
Key concepts
| Concept | What to understand |
|---|---|
| Relay IR | Functional IR; higher-level than LLVM IR |
| TE (Tensor Expressions) | Schedule-based computation description |
| AutoTVM / MetaSchedule | ML-guided kernel auto-tuning |
| Compilation pipeline | relay.optimize → te.create_prim_func → tvm.build |
Files
| File | Purpose |
|---|---|
lab.py | Stubs with TODO markers |
solution.py | Complete working implementation |
test_lab.py | pytest suite (skipped if TVM not installed) |
requirements.txt | Dependencies |
Run
pip install -r requirements.txt # TVM requires separate install: see tvm.apache.org
python lab.py
pytest test_lab.py -v