Skip to content

← Performance step 2 of 7

Medium Primitives

Four allocations, or one

((x + 1) * 2 - 3).clamp(min=0)

Four operations, four full-size tensors allocated, three of them discarded immediately. In eager mode there is no fusion: each kernel reads its input from memory, writes its output to memory, and the next one reads it straight back.

For an elementwise chain the arithmetic is trivial and the memory traffic is everything, so four passes over the data cost roughly four times one pass.

Reusing one buffer

out = torch.empty_like(x)
torch.add(x, 1, out=out)
torch.mul(out, 2, out=out)
torch.sub(out, 3, out=out)
out.clamp_(min=0)

One allocation. Each step reads and writes the same buffer, which also stays in cache between steps.

Be honest about what this costs you

That is four lines of noise replacing one clear expression, and it is only correct because nothing else refers to the intermediates. Inside a forward pass, in-place ops break autograd for the reasons the autograd track covered.

So this is a technique for inference and preprocessing, applied after measuring, in the one place that turned out to matter. Writing it everywhere makes a codebase worse.

What you should usually reach for instead

torch.compile(fn)

which fuses the whole chain into a single generated kernel: one pass, one allocation, and the source stays the readable expression. It is the right answer for elementwise chains in modern PyTorch, and this problem exists so you know what it is doing for you and can recognise when it has not.

Your task

def pipeline(x: torch.Tensor) -> torch.Tensor

Compute ((x + 1) * 2 - 3).clamp(min=0) allocating at most one tensor. x must be unchanged.