Skip to content

← Performance step 2 of 7

Medium Primitives

Reuse one output buffer

Jump to the task

This eager expression creates a result tensor at each of its four steps:

import torch

x = torch.tensor([0.0, 1.0, 2.0])
result = ((x + 1) * 2 - 3).clamp(min=0)
print(result.tolist())  # [0.0, 1.0, 3.0]

If no gradients are needed, you can reuse one destination instead of creating a separate result buffer for each step:

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)
print(out.tolist())  # [0.0, 1.0, 3.0]
print(x.tolist())    # [0.0, 1.0, 2.0]

Reuse is different from fusion

The second version still executes four arithmetic operations. It reuses storage but does not merge those operations into one kernel or remove all intermediate reads and writes. Cache behavior and execution time depend on the tensor size, device and implementation. Four operations do not imply four times the runtime of one operation.

Kernel fusion combines compatible operations so intermediate values need not be written to separate full-size tensors. A compiler such as torch.compile can sometimes fuse this kind of chain while keeping the original expression. The outcome depends on the backend, supported operations and graph capture; compilation is not a promise of one kernel or a speedup. Benchmark warmed-up execution and account separately for compilation cost.

Use this where gradients are not required

The out variants above reject autograd recording when an argument requires gradients. Some in-place operations do support autograd, but they can also invalidate values saved for backward. The rule is not that every in-place operation fails; the safe choice depends on which values backward needs. This exercise supplies inputs without gradients.

Keep the ordinary expression when clarity matters more than a measured benefit from buffer reuse. The exercise isolates storage behavior; it does not establish which implementation is fastest for your workload.

Your task

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

Compute ((x + 1) * 2 - 3).clamp(min=0) for a float32 vector, preserving x. Allocate one independent destination and reuse it for the remaining steps.

The harness expects an allocations count of exactly one. This field is a limited dispatch-level proxy: it counts tensor outputs whose storage pointer differs from the direct tensor arguments to that operation. It does not measure allocator calls, hidden kernel workspace, peak memory or runtime.

For x=[0.0, 1.0, 2.0], the harness reports {"values": [0.0, 1.0, 3.0], "allocations": 1, "x_unchanged": True}. Values are rounded to four decimals. Use the eager buffer-reuse pattern here; this exercise's counter is not a test of compiled kernel fusion.