Skip to content

Building C++ Extensions

Overview

A C++ extension is a compiled shared library that plugs into PyTorch as a new op — callable from Python, differentiable if you wrap it, and invisible to the Python overhead. This file covers the CPU path; the CUDA path is in Ch 09-02.

  • torch.utils.cpp_extension.load(...) — just-in-time compile from a single C++ string/file.
  • CppExtension / CUDAExtension in setup.py — proper packaging for distribution.
  • Signature: torch::Tensor in/out; use TORCH_CHECK for guards; at:: APIs mirror torch.
  • Register functions with PYBIND11_MODULE(TORCH_EXTENSION_NAME, m).

💡 Python's per-op overhead (~1-5µs) is exactly what a C++ op eliminates for tiny kernels.

Minimal JIT Build (from string)

#include <torch/extension.h>

torch::Tensor my_add(torch::Tensor a, torch::Tensor b) {
    TORCH_CHECK(a.sizes() == b.sizes(), "shape mismatch");
    return a + b;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("add", &my_add, "elementwise add");
}

from torch.utils.cpp_extension import load

ext = load(name="my_add_ext", sources=["my_add.cpp"], verbose=False)
print(ext.add(torch.tensor([1.0]), torch.tensor([2.0])))

// myops.cpp
#include <torch/extension.h>

torch::Tensor relu_twice(torch::Tensor x) {
    return torch::relu(x) + torch::relu(x);
}

torch::Tensor mul_tensor(torch::Tensor x, torch::Tensor y) {
    return x * y;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("relu_twice", &relu_twice);
    m.def("mul_tensor", &mul_tensor);
}

from torch.utils.cpp_extension import load

ext = load(name="myops", sources=["myops.cpp"],
           extra_cflags=["-O3"], extra_ldflags=[])
print(ext.relu_twice(torch.tensor([-1.0, 2.0])))

Packaging with setup.py (distributable)

# setup.py
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(
    name="myops",
    ext_modules=[
        CppExtension("myops", ["myops.cpp"], extra_compile_args=["-O3"]),
    ],
    cmdclass={"build_ext": BuildExtension},
)

pip install -e .        # builds and installs the extension
python -c "import myops; print(myops.relu_twice(torch.tensor([-1., 2.])))"

Autograd-Wrapping Your Op

import torch
import myops

class ReluTwiceFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return myops.relu_twice(x)
    @staticmethod
    def backward(ctx, g):
        # d/dx (relu+relu) = 2 if x>0 else 0  -- same as python relu chain
        return g * 2 * (x > 0).float()

💡 If your op is elementwise and trivial, torch already fuses it — only write C++ when python-vs-C++ speed matters (tiny ops, hot loops) or the op can't exist in Python.

Common Compile Pitfalls

Problem Fix
std::bad_alloc/crash TORCH_CHECK shapes/dtypes/contiguity first
Build slow cache with load(..., with_cuda=True, extra_include_paths=...); use -O3 -march=native
Missing symbol on load recompile after torch upgrade; rebuild from clean
JIT rebuild every time name it, reuse; force_reload=False
CPU-only env w/ CUDA code guard with with_cuda=False or separate paths

Testing Your Extension

import torch, myops

x = torch.randn(5, requires_grad=True)
out = myops.relu_twice(x)
out.sum().backward()
print("grad:", x.grad)

# compare vs pure python reference
assert torch.allclose(out.detach(), torch.relu(x) + torch.relu(x))
print("matches python reference ✓")

Key Takeaways

  • C++ extensions = compiled ops callable from Python; load() for JIT, setup.py for shipping.
  • Always TORCH_CHECK inputs; match return arity with torch::Tensor.
  • Wrap in autograd.Function for gradients; reference-check vs python.
  • Use them for tiny/hot ops and ops impossible in Python — not as a default.
  • CPU first; CUDA when profiling says GPU is idle waiting on CPU.