PyTorch ATen regions

A function whose body is entirely curated tensor operations compiles into one C++ region that calls ATen directly: one Python round trip per call instead of one per operator. layer is a matmul, an add, and a ReLU. The region runs them as PyTorch would, and autograd is unchanged.

Run it

python  torch_region.ppy
ppy     torch_region.ppy
ppy run torch_region.ppy

Through the dispatcher

@ppy.opt(3)
def layer(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
    return torch.relu(torch.add(torch.matmul(x, weight), bias))

The region does not reimplement matmul. Every at:: call inside it goes through PyTorch's dispatcher, so device selection, dtype promotion, and autograd behave as before: residual(tracked, tracked).sum().backward() fills tracked.grad. What the region removes is the Python interpreter between operators.

Built artifacts

A built artifact carries its regions. ppy build copies the extension beside the manifest, and the launcher loads it with no compiler in the process (torchrun).

Compared with PyTorch eager and torch.compile

The same layer on an 8×32 input, twenty thousand calls, in compare/: layer_bench.ppy, layer_eager.py, layer_compile.py. Milliseconds per call, to four places, best of five rounds, over five processes; one PyTorch thread.

• PPy is the function as written, and ppy run compiles it into one ATen region.
• PyTorch eager is the same function with no decorator: three dispatches from Python.
• torch.compile is the same function under the decorator, traced by Dynamo and written by Inductor.

@ppy.opt(3)
def layer(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
    return torch.relu(torch.add(torch.matmul(x, weight), bias))

@torch.compile
def layer(x, weight, bias):
    return torch.relu(torch.add(torch.matmul(x, weight), bias))

	PPy ATen region	PyTorch eager	torch.compile
layer, per call	0.0019 ± 0.0000	0.0020 ± 0.0000	0.0087 ± 0.0002

Three operators on a tensor this small cost about two microseconds either way. PyTorch's eager dispatch is cheap enough that the Python round trips the region removes are within the noise of the ATen calls themselves. The region's value is not this number. It is that the function keeps its source, its autograd, and its dispatcher, and that a built artifact carries it with no compiler in the process. torch.compile pays for its guards on every call, which on an 8×32 input is more than the work.

One layer of one shape is not the question anyone asks about torch.compile, and this table cannot answer it: the work is two microseconds, so every column is measuring its own overhead. 46_gpt2 asks it at model scale instead: GPT-2 XL, forty-eight blocks, one region each, on a datacenter GPU.

Intel Core Ultra 9 386H; PyTorch 2.14.0 (CPU) on CPython 3.13.13, PPy against the same PyTorch on CPython 3.14.5, from a checkout on a native filesystem.

Limitations

• A tensor subclass, a __torch_function__ override, or a device the region was not built for fails the guard, and the Python body runs.
• Building the region needs a C++ compiler and ninja. toolchain_ready() says what is missing.
• On an accelerator the region changes nothing measurable: kernel launch latency dominates the Python overhead it removes. So the program measures on the CPU, and on cuda too when one is present.

What it prints

python torch_region.ppy

torch 2.14.0+cpu | cuda False
# region active: False
layer on cpu              2.530 us/call   sample=596.816528
autograd survives the region: True

ppy torch_region.ppy

torch 2.14.0+cpu | cuda False
# region active: True
layer on cpu              2.069 us/call   sample=596.816528
autograd survives the region: True

ppy run torch_region.ppy

torch 2.14.0+cpu | cuda False
# region active: True
layer on cpu              2.118 us/call   sample=596.816528
autograd survives the region: True

Read on: Plugins: PyTorch · Training with torch

torch_region.ppy is hand-written; there is no .py source and no conversion step.

09_torch/torch_region.ppy

import time

import ppy
import torch


@ppy.opt(3)
def layer(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
    return torch.relu(torch.add(torch.matmul(x, weight), bias))


@ppy.opt(3)
def residual(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return torch.tanh(torch.add(x, torch.mul(y, 0.5)))


def report(label: str, elapsed: float, value: float) -> None:
    print(f"{label:22s} {elapsed * 1e6:8.3f} us/call   sample={value:.6f}")


def measure(name: str, x: torch.Tensor, w: torch.Tensor, b: torch.Tensor) -> None:
    rounds: int = 20000
    for _i in range(300):
        layer(x, w, b)
    if x.is_cuda:
        torch.cuda.synchronize()
    start: float = time.perf_counter()
    for _i in range(rounds):
        out = layer(x, w, b)
    if x.is_cuda:
        torch.cuda.synchronize()
    report(name, (time.perf_counter() - start) / rounds, float(out.sum()))


def main() -> None:
    torch.manual_seed(0)
    print("torch", torch.__version__, "| cuda", torch.cuda.is_available())
    print("# region active:", getattr(layer, "__ppy_region__", False))

    for device in ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]:
        x = torch.randn(8, 32, device=device)
        w = torch.randn(32, 32, device=device)
        b = torch.randn(32, device=device)
        measure(f"layer on {device}", x, w, b)

    tracked = torch.randn(4, 8, requires_grad=True)
    residual(tracked, tracked).sum().backward()
    print("autograd survives the region:", tracked.grad is not None)


if __name__ == "__main__":
    main()

Counterpart programs

The programs the comparison above measured, each written the way its tool expects. The PPy one is first.

layer_bench.ppy (PPy)

"""`layer` on an 8x32 CPU input, twenty thousand calls: one ATen region per call. The PPY side."""

import time

import ppy
import torch


@ppy.opt(3)
def layer(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
    return torch.relu(torch.add(torch.matmul(x, weight), bias))


def main() -> None:
    torch.manual_seed(0)
    torch.set_num_threads(1)
    x = torch.randn(8, 32)
    w = torch.randn(32, 32)
    b = torch.randn(32)
    for _i in range(300):
        layer(x, w, b)
    rounds: int = 20000
    best = 1e9
    for _ in range(5):
        start: float = time.perf_counter()
        for _i in range(rounds):
            out = layer(x, w, b)
        best = min(best, (time.perf_counter() - start) / rounds * 1000.0)
    print(f"# layer, per call: {best:.4f} ms")
    print(f"{float(out.sum()):.4f}")


main()

layer_compile.py (Python)

"""The same function under `torch.compile`: Dynamo traces it, Inductor writes a fused kernel."""

import time

import torch


@torch.compile
def layer(x, weight, bias):
    return torch.relu(torch.add(torch.matmul(x, weight), bias))


def main():
    torch.manual_seed(0)
    torch.set_num_threads(1)
    x = torch.randn(8, 32)
    w = torch.randn(32, 32)
    b = torch.randn(32)
    for _ in range(300):
        layer(x, w, b)
    rounds = 20000
    best = 1e9
    for _ in range(5):
        start = time.perf_counter()
        for _ in range(rounds):
            out = layer(x, w, b)
        best = min(best, (time.perf_counter() - start) / rounds * 1000.0)
    print(f"# layer, per call: {best:.4f} ms")
    print(f"{float(out.sum()):.4f}")


main()

layer_eager.py (Python)

"""The same three operators as PyTorch runs them eagerly: one Python round trip per operator."""

import time

import torch


def layer(x, weight, bias):
    return torch.relu(torch.add(torch.matmul(x, weight), bias))


def main():
    torch.manual_seed(0)
    torch.set_num_threads(1)
    x = torch.randn(8, 32)
    w = torch.randn(32, 32)
    b = torch.randn(32)
    for _ in range(300):
        layer(x, w, b)
    rounds = 20000
    best = 1e9
    for _ in range(5):
        start = time.perf_counter()
        for _ in range(rounds):
            out = layer(x, w, b)
        best = min(best, (time.perf_counter() - start) / rounds * 1000.0)
    print(f"# layer, per call: {best:.4f} ms")
    print(f"{float(out.sum()):.4f}")


main()

Source: examples/09_torch.