PyTorch C++ Ops Research¶

What you'll learn¶

• Why write custom PyTorch operators
• The C++ extension mechanism
• Registering a custom op
• Custom autograd (forward + backward)
• When it's worth it

Sometimes PyTorch's built-in operations aren't enough — you need a custom operator for a novel computation or maximum speed. PyTorch lets you write ops in C++/CUDA and call them from Python as if native. This extends the general C++ Extensions topic to PyTorch specifically.

Why custom operators¶

A core question explored in PyTorch C++ Ops: Why custom operators.

Reasons to drop from Python into C++/CUDA for a PyTorch op:

• Speed — a fused custom kernel avoids the overhead of composing many small PyTorch ops (each of which launches a GPU kernel and moves memory).
• Novel operations — implement something PyTorch doesn't provide (a new attention variant, a custom loss).
• Kernel fusion — combine several operations into one GPU pass, avoiding intermediate memory traffic (like GPU Kernels).

The pattern is the familiar one: Python for the model, C++/CUDA for the hot custom op.

The extension mechanism¶

Compile C++/CUDA and bind it as a PyTorch operator.

PyTorch exposes its tensor library (ATen) to C++, so your operator works with PyTorch tensors directly:

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

torch::Tensor custom_add(torch::Tensor a, torch::Tensor b) {
    return a + b;                      // uses PyTorch's tensor ops in C++
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("custom_add", &custom_add, "A custom add");
}

import torch
from torch.utils.cpp_extension import load

# JIT-compile the C++ on first import
ext = load(name="ext", sources=["custom_op.cpp"])
result = ext.custom_add(torch.tensor([1.0]), torch.tensor([2.0]))   # -> tensor([3.])

PyTorch can compile the C++ just-in-time (load) or ahead-of-time via setuptools. It builds on pybind11 (see C++ Extensions) plus PyTorch's tensor types.

Custom autograd (forward + backward)¶

Define both passes so your op works with autograd.

The key ML-specific piece: for your op to work in training, it must support backpropagation — you provide both the forward computation and the gradient (backward). This is autograd applied at the operator level:

import torch

class CustomFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)       # remember inputs for backward
        return x * x                    # forward: y = x²

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        return grad_output * 2 * x      # backward: dy/dx = 2x

You define forward (the computation) and backward (its gradient) — exactly the pattern from the autograd engine we built, but registered with PyTorch so it composes with the rest of the network. Get the backward math wrong and training silently breaks — so you check gradients against numerical differentiation (torch.autograd.gradcheck).

When it's worth it¶

A core question explored in PyTorch C++ Ops: When it's worth it.

Practice exercises¶

1. Explain why fusing several PyTorch ops into one custom kernel can be faster.
2. Describe why a custom op needs both forward and backward to be usable in training.
3. Relate PyTorch's Function.forward/backward to the autograd engine from the Scientific section.
4. Explain what gradcheck does and why you'd run it on a custom op.
5. Order these by effort for a custom fast op: built-in ops, Triton kernel, C++/CUDA op — and say when each is justified.