Lab 02 — Custom torch.compile Backend

Phase: 04 — torch.compile, FX Graph & Custom Backends | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

Implement a custom backend that classifies ops as NPU-eligible or CPU-fallback and substitutes them in the FX graph.

What you build

  • NPU op registry with @register_npu_op decorator
  • classify_nodes — analyze FX graph for NPU vs CPU ops
  • transform_graph — replace ATen ops with registered NPU ops via FX node substitution
  • npu_backend — torch.compile(model, backend=npu_backend) callable
  • verify_backend_correctness — compare eager vs compiled numerics

Key concepts

ConceptWhat to understand
Backend callable(gm: GraphModule, example_inputs) → Callable
ATen opsaten.mm.default, aten.relu.default — canonical op names
Node replacementgraph.inserting_after(node) + node.replace_all_uses_with
Correctness gateMax absolute diff < 1e-4 between eager and compiled

Files

FilePurpose
lab.pyStubs with TODO markers
solution.pyComplete working implementation
test_lab.pypytest suite
requirements.txtDependencies

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v