Lab 02 — Custom C++ Operator

Phase: 01 — PyTorch & TF Framework Internals | Difficulty: ⭐⭐⭐⭐☆ | Time: 2–4 hours

Read ../HITCHHIKERS-GUIDE.md before starting.

What you build

A custom PyTorch operator written in C++ that registers with the dispatcher:

  • C++ extension using pybind11 and torch/extension.h
  • Forward pass: fused linear + ReLU in a single C++ kernel
  • Registration via TORCH_LIBRARY / PYBIND11_MODULE
  • Python test verifying numeric equality with the PyTorch reference

Key concepts

ConceptWhat to understand
ATen dispatcherHow torch.ops.mylib.op() routes to C++
at::TensorC++ tensor API — same data, different language
setup.py + torch.utils.cpp_extensionBuild system for extensions
SchemaThe type signature registered with the dispatcher

Files

FilePurpose
setup.pyBuild configuration for C++ extension
csrc/C++ source files
solution.pyPython tests and usage
requirements.txtPython dependencies

Run

pip install -r requirements.txt
python setup.py build_ext --inplace
python solution.py